fix: PEFT parameter statistics only count the first model chunk

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>

Author: andrewwhitecdw

src/megatron/bridge/training/setup.py

@@ -613,14 +613,15 @@ def _apply_peft_transformation(peft, base_model: list[MegatronModule]) -> list[M
     peft.set_params_to_save(transformed_model)
 
     # Log PEFT statistics
-    model_to_analyze = transformed_model[0] if isinstance(transformed_model, list) else transformed_model
+    model_chunks = transformed_model if isinstance(transformed_model, list) else [transformed_model]
     total_params = 0
     trainable_params = 0
-    for param in model_to_analyze.parameters():
-        param_count = param.numel()
-        total_params += param_count
-        if param.requires_grad:
-            trainable_params += param_count
+    for model_chunk in model_chunks:
+        for param in model_chunk.parameters():
+            param_count = param.numel()
+            total_params += param_count
+            if param.requires_grad:
+                trainable_params += param_count
 
     print_rank_0("PEFT Statistics:")
     print_rank_0(f"  Total parameters: {total_params:,}")

tests/unit_tests/bridge/training/test_setup.py

@@ -0,0 +1,43 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from megatron.bridge.training.setup import _apply_peft_transformation
+
+
+class _DummyModelChunk(nn.Module):
+    def __init__(self, total, trainable):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(total), requires_grad=trainable)
+
+
+class _DummyPEFT:
+    def __call__(self, base_model, training=True):
+        return base_model
+
+    def set_params_to_save(self, model):
+        pass
+
+
+def test_apply_peft_transformation_counts_all_chunks():
+    """PEFT parameter stats must account for every model chunk, not only chunk 0."""
+    chunk0 = _DummyModelChunk(10, False)
+    chunk1 = _DummyModelChunk(5, True)
+    peft = _DummyPEFT()
+
+    logs = []
+
+    def fake_print_rank_0(msg):
+        logs.append(msg)
+
+    import megatron.bridge.training.setup as setup_module
+
+    orig_print_rank_0 = setup_module.print_rank_0
+    setup_module.print_rank_0 = fake_print_rank_0
+    try:
+        _apply_peft_transformation(peft, [chunk0, chunk1])
+    finally:
+        setup_module.print_rank_0 = orig_print_rank_0
+
+    assert any("Total parameters: 15" in msg for msg in logs)
+    assert any("Trainable parameters: 5" in msg for msg in logs)