[michaellchung]

Summary

• Learner: enter training mode before enabling gradient checkpointing. from_pretrained returns an eval-mode module and the learner never called model.train(). Transformers' GradientCheckpointingLayer only recomputes when self.training is true, so --gradient-checkpointing on was accepted and logged but silently ineffective, and dropout (LoRA's included) was inert. Setup now lives in enter_training_mode, with the ordering covered by tests. On one Ascend 910B3 with Qwen3-0.6B LoRA, activation peak at 16k tokens drops from 59.0 GB to 10.1 GB, and 28k tokens goes from OOM to 10.9 GB.
• accel: recognize an allocator OOM without a vendor exception class. torch.npu exports no OutOfMemoryError, so Ascend's bare RuntimeError("NPU out of memory ...") escaped both memory probes (causal-LM auto-batch and diffusion). accel.is_oom(exc, device) matches the vendor class, torch's generic OutOfMemoryError and the message-only form; anything else is re-raised.
• Autobatch: probe the real sft_loss. The stand-in logits[:, :-1].float().sum() ignored the log-softmax tensor that F.cross_entropy keeps for backward, so the probe admitted a micro-batch about one f32 logits tensor too large. The probe now runs the native loss with all-ones weights.
• Learner: name the chunked loss kernel in the recipe log. Each backend is mapped explicitly and unknown values are rejected.
• Learner: keep first-step warm-up out of the s/step telemetry. The first logging window included the one-off first-step cost (dataset build, compilation, collective setup); the first step's duration is now logged on its own and the first window starts after it.

Tests

• tests/test_learner_units.py, tests/test_accel.py, tests/test_autobatch.py, tests/test_chunked_ce.py.

🤖 Generated with Claude Code