Adds scripts/10_train.py, the Phase 3 LoRA training script for the Qwen3.6-35B-A3B base checkpoint.
Uses the Phase 2 chat template (data/chat_template_train.jinja) with manual label masking via tokenizer.apply_chat_template(..., return_assistant_tokens_mask=True): labels equal input_ids where assistant_masks == 1, -100 everywhere else, so the model never trains on system/user/tool tokens.
Loads the base model with AutoModelForCausalLM.from_pretrained(..., dtype=torch.bfloat16, attn_implementation="flash_attention_2") (no vision tower instantiated).
Configures LoRA (r=32, lora_alpha=64, lora_dropout=0.05, bias="none") over the confirmed target modules: q/k/v/o_proj (full attention), in_proj_qkv/z/a/b + out_proj (Gated DeltaNet linear attention), and shared_expert.gate/up/down_proj (MoE).
Adds a --max-steps flag / MAX_STEPS env var for dry-run support, and logs torch.cuda.max_memory_allocated() (actual peak memory) at the end of training.
Changes
File
Change
scripts/10_train.py
New Phase 3 LoRA training script (chat template masking, model + LoRA config, Trainer setup, dry-run support, peak memory logging)
Validation
All runs executed inside the qwen-lora-train container on the spark server via docker exec (no local GPU available).
Dry-run, 20 steps: no exceptions, 0.122% trainable params (expected for this LoRA config), loss stable in the 0.32-0.37 range, eval_loss 0.349, peak memory 103.68GB (of ~110GB available, confirmed after checking vllm-qwen36 was stopped). Full run was explicitly approved by the user after reviewing these results.
Full run, 2 epochs / 166 steps: completed without errors. Final train loss 0.2735, final eval_loss 0.2761 (down from ~0.35 at the start). Wall clock time 6516s (~1h49m). Peak memory identical to the dry-run (103.68GB, no degradation). Checkpoints saved to out/lora-adapter/ every 50 steps (checkpoint-100, checkpoint-150, checkpoint-166; save_total_limit=3 pruned checkpoint-50). The final adapter (adapter_model.safetensors, ~169MB) and tokenizer are saved there but are not committed (binary training artifacts, covered by .gitignore).
Non-blocking note: torchao was upgraded (0.15.0 -> 0.17.0) inside the training container because peft 0.19.1 raises an unconditional exception in its LoRA dispatcher when it detects an old torchao version, even though this training does not use torchao quantization. This was an environment fix inside the container, not a code change in this repo.
Safety: docker stop/docker start was never run against vllm-qwen36 — only read-only docker ps checks before each run, both times confirming it was stopped. It remains stopped; restarting it is the user's responsibility.
qa-validator was intentionally skipped: this is a pure training task with no web-facing interface to validate.
Test Plan
Dry-run of 20 training steps completed without exceptions inside the training container
Full 2-epoch (166-step) training run completed without errors; train/eval loss decreased consistently
Verified peak GPU memory stayed within budget (103.68GB) across both runs
Confirmed vllm-qwen36 container state was only read, never stopped/started
## Summary
- Adds `scripts/10_train.py`, the Phase 3 LoRA training script for the Qwen3.6-35B-A3B base checkpoint.
- Uses the Phase 2 chat template (`data/chat_template_train.jinja`) with manual label masking via `tokenizer.apply_chat_template(..., return_assistant_tokens_mask=True)`: labels equal `input_ids` where `assistant_masks == 1`, `-100` everywhere else, so the model never trains on system/user/tool tokens.
- Loads the base model with `AutoModelForCausalLM.from_pretrained(..., dtype=torch.bfloat16, attn_implementation="flash_attention_2")` (no vision tower instantiated).
- Configures LoRA (`r=32`, `lora_alpha=64`, `lora_dropout=0.05`, `bias="none"`) over the confirmed target modules: `q/k/v/o_proj` (full attention), `in_proj_qkv/z/a/b` + `out_proj` (Gated DeltaNet linear attention), and `shared_expert.gate/up/down_proj` (MoE).
- Uses `transformers.Trainer` (not TRL `SFTTrainer`) since label masking is already handled manually. `TrainingArguments`: bf16, gradient checkpointing, `optim=adamw_8bit`, `lr=1e-4` cosine schedule, `warmup_ratio=0.03`, `num_train_epochs=2`, `per_device_train_batch_size=1`, `gradient_accumulation_steps=16` (effective batch 16), eval/save every 50 steps, `save_total_limit=3`.
- Adds a `--max-steps` flag / `MAX_STEPS` env var for dry-run support, and logs `torch.cuda.max_memory_allocated()` (actual peak memory) at the end of training.
## Changes
| File | Change |
|------|--------|
| `scripts/10_train.py` | New Phase 3 LoRA training script (chat template masking, model + LoRA config, Trainer setup, dry-run support, peak memory logging) |
## Validation
All runs executed inside the `qwen-lora-train` container on the `spark` server via `docker exec` (no local GPU available).
1. **Dry-run, 20 steps**: no exceptions, 0.122% trainable params (expected for this LoRA config), loss stable in the 0.32-0.37 range, eval_loss 0.349, peak memory 103.68GB (of ~110GB available, confirmed after checking `vllm-qwen36` was stopped). Full run was explicitly approved by the user after reviewing these results.
2. **Full run, 2 epochs / 166 steps**: completed without errors. Final train loss 0.2735, final eval_loss 0.2761 (down from ~0.35 at the start). Wall clock time 6516s (~1h49m). Peak memory identical to the dry-run (103.68GB, no degradation). Checkpoints saved to `out/lora-adapter/` every 50 steps (`checkpoint-100`, `checkpoint-150`, `checkpoint-166`; `save_total_limit=3` pruned `checkpoint-50`). The final adapter (`adapter_model.safetensors`, ~169MB) and tokenizer are saved there but are not committed (binary training artifacts, covered by `.gitignore`).
Non-blocking note: `torchao` was upgraded (0.15.0 -> 0.17.0) inside the training container because `peft 0.19.1` raises an unconditional exception in its LoRA dispatcher when it detects an old `torchao` version, even though this training does not use `torchao` quantization. This was an environment fix inside the container, not a code change in this repo.
Safety: `docker stop`/`docker start` was never run against `vllm-qwen36` — only read-only `docker ps` checks before each run, both times confirming it was stopped. It remains stopped; restarting it is the user's responsibility.
`qa-validator` was intentionally skipped: this is a pure training task with no web-facing interface to validate.
## Test Plan
- [x] Dry-run of 20 training steps completed without exceptions inside the training container
- [x] Full 2-epoch (166-step) training run completed without errors; train/eval loss decreased consistently
- [x] Verified peak GPU memory stayed within budget (103.68GB) across both runs
- [x] Confirmed `vllm-qwen36` container state was only read, never stopped/started
Blocking a user prevents them from interacting with repositories, such as opening or commenting on pull requests or issues. Learn more about blocking a user.
Summary
scripts/10_train.py, the Phase 3 LoRA training script for the Qwen3.6-35B-A3B base checkpoint.data/chat_template_train.jinja) with manual label masking viatokenizer.apply_chat_template(..., return_assistant_tokens_mask=True): labels equalinput_idswhereassistant_masks == 1,-100everywhere else, so the model never trains on system/user/tool tokens.AutoModelForCausalLM.from_pretrained(..., dtype=torch.bfloat16, attn_implementation="flash_attention_2")(no vision tower instantiated).r=32,lora_alpha=64,lora_dropout=0.05,bias="none") over the confirmed target modules:q/k/v/o_proj(full attention),in_proj_qkv/z/a/b+out_proj(Gated DeltaNet linear attention), andshared_expert.gate/up/down_proj(MoE).transformers.Trainer(not TRLSFTTrainer) since label masking is already handled manually.TrainingArguments: bf16, gradient checkpointing,optim=adamw_8bit,lr=1e-4cosine schedule,warmup_ratio=0.03,num_train_epochs=2,per_device_train_batch_size=1,gradient_accumulation_steps=16(effective batch 16), eval/save every 50 steps,save_total_limit=3.--max-stepsflag /MAX_STEPSenv var for dry-run support, and logstorch.cuda.max_memory_allocated()(actual peak memory) at the end of training.Changes
scripts/10_train.pyValidation
All runs executed inside the
qwen-lora-traincontainer on thesparkserver viadocker exec(no local GPU available).vllm-qwen36was stopped). Full run was explicitly approved by the user after reviewing these results.out/lora-adapter/every 50 steps (checkpoint-100,checkpoint-150,checkpoint-166;save_total_limit=3prunedcheckpoint-50). The final adapter (adapter_model.safetensors, ~169MB) and tokenizer are saved there but are not committed (binary training artifacts, covered by.gitignore).Non-blocking note:
torchaowas upgraded (0.15.0 -> 0.17.0) inside the training container becausepeft 0.19.1raises an unconditional exception in its LoRA dispatcher when it detects an oldtorchaoversion, even though this training does not usetorchaoquantization. This was an environment fix inside the container, not a code change in this repo.Safety:
docker stop/docker startwas never run againstvllm-qwen36— only read-onlydocker pschecks before each run, both times confirming it was stopped. It remains stopped; restarting it is the user's responsibility.qa-validatorwas intentionally skipped: this is a pure training task with no web-facing interface to validate.Test Plan
vllm-qwen36container state was only read, never stopped/started