[Bug] scripts/chat_sft.py produces loss: nan from step 00001 on small device-batch-size (≤8) due to fully-masked micro-batches

Description:

When running Supervised Fine-Tuning (SFT) on hardware with limited VRAM (e.g., 6GB/8GB GPUs), setting a small device-batch-size (e.g., 4 or 8) consistently leads to NaN loss values starting from Step 1. This appears to be a regression or unhandled edge case in how the SFT trainer handles micro-batches that contain no active training targets.

Log

(nanochat) ss@ss-Predator-PH315-53:~/nanochat$ WANDB_RUN=nano_d12_sft_ss_v_stable_final \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
python -m scripts.chat_sft \
  --model-tag nanochat_v1_d12 \
  --model-step 2205 \
  --device-batch-size 8 \
  --matrix-lr 0.0002 \
  --warmup-ratio 0.05 \
  --run nano_d12_sft_ss_v_stable_final
Autodetected device type: cuda
2026-03-05 21:46:46,861 - nanochat.common - INFO - Distributed world size: 1
COMPUTE_DTYPE: torch.bfloat16 (auto-detected: CUDA SM 86 (bf16 supported))
2026-03-05 21:46:46,861 - nanochat.common - WARNING - Peak flops undefined for: NVIDIA GeForce RTX 3060 Laptop GPU, MFU will show as 0%
GPU: NVIDIA GeForce RTX 3060 Laptop GPU | Peak FLOPS (BF16): inf
wandb: Currently logged in as: suraj03143 (suraj-self) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
wandb: Tracking run with wandb version 0.21.3
wandb: Run data is saved locally in /home/ss/nanochat/wandb/run-20260305_214648-ehs5d2d2
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run nano_d12_sft_ss_v_stable_final
wandb: ⭐️ View project at https://wandb.ai/suraj-self/nanochat-sft
wandb:  View run at https://wandb.ai/suraj-self/nanochat-sft/runs/ehs5d2d2
WARNING: Flash Attention 3 not available, using PyTorch SDPA fallback. Training will be less efficient.
2026-03-05 21:46:51,234 - nanochat.checkpoint_manager - INFO - Loading model from /home/ss/.cache/nanochat/base_checkpoints/nanochat_v1_d12 with step 2205
2026-03-05 21:46:52,112 - nanochat.checkpoint_manager - INFO - Building model with config: {'sequence_len': 512, 'vocab_size': 32768, 'n_layer': 12, 'n_head': 6, 'n_kv_head': 6, 'n_embd': 768, 'window_pattern': 'SSSL'}
Inherited max_seq_len=512 from pretrained checkpoint
NOTE: --device-batch-size=8 overrides pretrained value of 4
Inherited total_batch_size=524288 from pretrained checkpoint
Inherited embedding_lr=0.3 from pretrained checkpoint
Inherited unembedding_lr=0.004 from pretrained checkpoint
NOTE: --matrix-lr=0.0002 overrides pretrained value of 0.02
Tokens / micro-batch / rank: 8 x 512 = 4,096
Tokens / micro-batch: 4,096
Total batch size 524,288 => gradient accumulation steps: 128
Scaling the LR for the AdamW parameters ∝1/√(768/768) = 1.000000
2026-03-05 21:46:53,404 - nanochat.checkpoint_manager - INFO - Loading optimizer state from /home/ss/.cache/nanochat/base_checkpoints/nanochat_v1_d12/optim_002205_rank0.pt
Loaded optimizer state from pretrained checkpoint (momentum buffers only, LRs reset)
Training mixture: 1,071,759 rows (MMLU x3, GSM8K x4)
W0305 21:47:34.063000 19137 .venv/lib/python3.10/site-packages/torch/_inductor/utils.py:1613] [0/0] Not enough SMs to use max_autotune_gemm mode
Step 00000 | Validation bpb: 1.0855
step 00001 (0.02%) | loss: nan | lrm: 0.00 | dt: 64613.10ms | tok/sec: 8,114 | mfu: 0.00 | epoch: 1 | total time: 0.00m
step 00002 (0.02%) | loss: nan | lrm: 0.00 | dt: 25134.97ms | tok/sec: 20,858 | mfu: 0.00 | epoch: 1 | total time: 0.00m
step 00003 (0.02%) | loss: nan | lrm: 0.00 | dt: 25183.07ms | tok/sec: 20,819 | mfu: 0.00 | epoch: 1 | total time: 0.00m

The Root Cause:

In SFT, we mask the "User" portion of the conversation (setting targets to -1) and only calculate loss on the "Assistant" response.

• When the device-batch-size is small, there is a statistically high probability that a micro-batch will contain only "User" tokens or padding (especially at the start of long conversations).
• torch.nn.functional.cross_entropy with reduction='mean' returns NaN when all labels are -1 because the denominator (number of active tokens) is zero.
• Once a single micro-batch in a gradient accumulation cycle returns NaN, the global gradient becomes NaN, poisoning the model weights permanently.

Hardware Environment:

• GPU: NVIDIA RTX 3060 Laptop (6GB VRAM)
• Memory: 16GB RAM
• PyTorch: 2.9.1+cu128
• Context: Using small batch sizes to avoid OOM on consumer hardware.

Steps to Reproduce:

1. Use a pretrained d12 model.
2. Run the SFT script with a micro-batch size of 4 or 8

python -m scripts.chat_sft \
  --model-tag [YOUR_MODEL_TAG] \
  --model-step [YOUR_STEP] \
  --device-batch-size 8 \
  --run sft_reproduction_test

1. Observe loss: nan from Step 1

Code Review:

The issue occurs in the training loop where the backward pass is called on the loss. If the loss is NaN, the gradients are corrupted.

Suggested Resolution

The trainer should check if a micro-batch contains any valid targets before computing loss or performing a backward pass. A safety check such as if (y != -1).any(): should be implemented. If no targets are present, the micro-batch should contribute zero to the gradient accumulation rather than NaN