Merge 55eb345515 into 7808dc7159

scripts/chat_sft.py

@@ -334,6 +334,8 @@ smooth_train_loss = 0 # EMA of training loss
 ema_beta = 0.9 # EMA decay factor
 total_training_time = 0 # total wall-clock time of training
 step = 0
+skipped_batches = 0 # count of micro-batches skipped due to no valid targets
+warned_about_skips = False # whether we've warned the user about skipped batches
 while True:
     flops_so_far = num_flops_per_token * args.total_batch_size * step
 
@@ -430,13 +432,33 @@ while True:
     synchronize()
     t0 = time.time()
     for micro_step in range(grad_accum_steps):
-        loss = model(x, y)
-        train_loss = loss.detach() # for logging
-        loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here
-        if scaler is not None:
-            scaler.scale(loss).backward()
+        # Check if this micro-batch has any valid (non-masked) targets
+        # In SFT, targets are -1 for "User" portions; loss is only computed on "Assistant" responses
+        # If all targets are -1, cross_entropy with reduction='mean' returns NaN (division by zero)
+        has_valid_targets = (y != -1).any()
+        
+        if has_valid_targets:
+            loss = model(x, y)
+            train_loss = loss.detach() # for logging
+            loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here
+            if scaler is not None:
+                scaler.scale(loss).backward()
+            else:
+                loss.backward()
         else:
-            loss.backward()
+            # Skip this micro-batch: no valid targets, would cause NaN
+            # Set train_loss to 0 for logging (this micro-batch contributes nothing)
+            train_loss = torch.tensor(0.0, device=device)
+            # Note: We don't call .backward() here, so this micro-batch contributes zero gradient
+            # This is correct: there's nothing to learn from a fully-masked batch
+            skipped_batches += 1
+            # Warn user once if skipping becomes frequent (>5% of batches in first 100 steps)
+            if not warned_about_skips and step < 100 and skipped_batches > (step * grad_accum_steps * 0.05):
+                print0(f"WARNING: Skipping micro-batches with no valid targets ({skipped_batches} so far).")
+                print0(f"  This is normal with small device-batch-size, but if it happens frequently,")
+                print0(f"  consider increasing --device-batch-size to reduce wasted computation.")
+                warned_about_skips = True
+        
         x, y = next(train_loader) # prefetch the next batch while the GPU is busy with forward/backward
         progress = max(progress, approx_progress) # only increase progress monotonically
     # step the optimizer

tests/test_sft_masked_batches.py

@@ -0,0 +1,161 @@
+"""
+Test SFT training with fully-masked micro-batches.
+This test verifies the fix for issue #590: NaN loss when device-batch-size is small.
+
+Run with:
+python -m pytest tests/test_sft_masked_batches.py -v
+"""
+
+import torch
+import torch.nn.functional as F
+from nanochat.gpt import GPT, GPTConfig
+
+
+def test_fully_masked_batch_no_nan():
+    """
+    Test that a fully-masked batch (all targets = -1) doesn't cause NaN loss.
+    
+    Before the fix (#590), this would return NaN because cross_entropy with
+    reduction='mean' divides by zero when all targets are ignored.
+    """
+    # Create a minimal model
+    config = GPTConfig(
+        sequence_len=16,
+        vocab_size=100,
+        n_layer=2,
+        n_head=2,
+        n_kv_head=2,
+        n_embd=32,
+        window_pattern="L"
+    )
+    model = GPT(config)
+    model.eval()  # disable dropout for deterministic testing
+    
+    # Create a batch where ALL targets are masked (-1)
+    batch_size = 4
+    seq_len = 16
+    inputs = torch.randint(0, config.vocab_size, (batch_size, seq_len))
+    targets = torch.full((batch_size, seq_len), -1, dtype=torch.long)  # all masked
+    
+    # Forward pass - should NOT return NaN
+    with torch.no_grad():
+        loss = model(inputs, targets, loss_reduction='mean')
+    
+    # The model's forward pass uses cross_entropy with ignore_index=-1
+    # With all targets masked, the denominator is 0, causing NaN
+    # This test documents the current behavior
+    assert torch.isnan(loss), (
+        "Model forward pass should return NaN for fully-masked batch. "
+        "The training loop (chat_sft.py) must detect and skip such batches."
+    )
+
+
+def test_partially_masked_batch_valid_loss():
+    """
+    Test that a partially-masked batch (some targets = -1) returns valid loss.
+    """
+    config = GPTConfig(
+        sequence_len=16,
+        vocab_size=100,
+        n_layer=2,
+        n_head=2,
+        n_kv_head=2,
+        n_embd=32,
+        window_pattern="L"
+    )
+    model = GPT(config)
+    model.eval()
+    
+    # Create a batch where SOME targets are masked
+    batch_size = 4
+    seq_len = 16
+    inputs = torch.randint(0, config.vocab_size, (batch_size, seq_len))
+    targets = torch.randint(0, config.vocab_size, (batch_size, seq_len))
+    
+    # Mask out first half of sequence (simulating "User" portion in SFT)
+    targets[:, :seq_len//2] = -1
+    
+    # Forward pass - should return valid (non-NaN) loss
+    with torch.no_grad():
+        loss = model(inputs, targets, loss_reduction='mean')
+    
+    assert not torch.isnan(loss), "Loss should be valid for partially-masked batch"
+    assert loss.item() > 0, "Loss should be positive"
+
+
+def test_sft_batch_validation_logic():
+    """
+    Test the validation logic that should be used in chat_sft.py training loop.
+    
+    This simulates the fix: check (y != -1).any() before computing loss.
+    """
+    batch_size = 4
+    seq_len = 16
+    
+    # Test case 1: Fully masked batch
+    targets_fully_masked = torch.full((batch_size, seq_len), -1, dtype=torch.long)
+    has_valid_targets = (targets_fully_masked != -1).any()
+    assert not has_valid_targets, "Fully masked batch should have no valid targets"
+    
+    # Test case 2: Partially masked batch
+    targets_partial = torch.randint(0, 100, (batch_size, seq_len))
+    targets_partial[:, :seq_len//2] = -1
+    has_valid_targets = (targets_partial != -1).any()
+    assert has_valid_targets, "Partially masked batch should have valid targets"
+    
+    # Test case 3: No masking
+    targets_unmasked = torch.randint(0, 100, (batch_size, seq_len))
+    has_valid_targets = (targets_unmasked != -1).any()
+    assert has_valid_targets, "Unmasked batch should have valid targets"
+
+
+def test_cross_entropy_behavior_with_ignore_index():
+    """
+    Document the cross_entropy behavior that causes the NaN issue.
+    
+    This test shows why the fix is necessary at the training loop level.
+    """
+    # Setup
+    vocab_size = 100
+    batch_size = 4
+    seq_len = 16
+    
+    # Create random logits
+    logits = torch.randn(batch_size * seq_len, vocab_size)
+    
+    # Test 1: All targets masked -> NaN with reduction='mean'
+    targets_all_masked = torch.full((batch_size * seq_len,), -1, dtype=torch.long)
+    loss_mean = F.cross_entropy(logits, targets_all_masked, ignore_index=-1, reduction='mean')
+    assert torch.isnan(loss_mean), "cross_entropy should return NaN when all targets are ignored"
+    
+    # Test 2: All targets masked -> 0.0 with reduction='sum'
+    loss_sum = F.cross_entropy(logits, targets_all_masked, ignore_index=-1, reduction='sum')
+    assert loss_sum.item() == 0.0, "cross_entropy with reduction='sum' should return 0 when all targets ignored"
+    
+    # Test 3: Some targets valid -> valid loss with reduction='mean'
+    targets_partial = torch.randint(0, vocab_size, (batch_size * seq_len,))
+    targets_partial[:batch_size * seq_len // 2] = -1  # mask half
+    loss_partial = F.cross_entropy(logits, targets_partial, ignore_index=-1, reduction='mean')
+    assert not torch.isnan(loss_partial), "cross_entropy should return valid loss for partially masked targets"
+    assert loss_partial.item() > 0, "Loss should be positive"
+
+
+if __name__ == "__main__":
+    # Run tests manually
+    print("Running test_fully_masked_batch_no_nan...")
+    test_fully_masked_batch_no_nan()
+    print("✓ Passed")
+    
+    print("Running test_partially_masked_batch_valid_loss...")
+    test_partially_masked_batch_valid_loss()
+    print("✓ Passed")
+    
+    print("Running test_sft_batch_validation_logic...")
+    test_sft_batch_validation_logic()
+    print("✓ Passed")
+    
+    print("Running test_cross_entropy_behavior_with_ignore_index...")
+    test_cross_entropy_behavior_with_ignore_index()
+    print("✓ Passed")
+    
+    print("\nAll tests passed!")