fix: PEFT parameter statistics only count the first model chunk

[andrewwhitecdw]

This PR addresses the following issue in src/megatron/bridge/training/setup.py: PEFT parameter statistics only count the first model chunk.

Changes

• src/megatron/bridge/training/setup.py: PEFT parameter statistics only count the first model chunk.

Details

--- a/src/megatron/bridge/training/setup.py
+++ b/src/megatron/bridge/training/setup.py
@@ -1,14 +1,16 @@
-    # Log PEFT statistics
-    model_to_analyze = transformed_model[0] 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
-
-    print_rank_0("PEFT Statistics:")
-    print_rank_0(f"  Total parameters: {total_params:,}")
-    print_rank_0(f"  Trainable parameters: {trainable_params:,}")
-    print_rank_0(f"  Trainable percentage: {100 * trainable_params / total_params:.2f}%")
+    # Log PEFT statistics
+    model_chunks = transformed_model if isinstance(transformed_model, list) else [transformed_model]
+    total_params = 0
+    trainable_params = 0
+    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:,}")
+    print_rank_0(f"  Trainable parameters: {trainable_params:,}")
+    trainable_pct = 100 * trainable_params / total_params if total_params > 0 else 0.0
+    print_rank_0(f"  Trainable percentage: {trainable_pct:.2f}%")

Tests

• tests/unit_tests/bridge/training/test_setup.py

diff --git a/tests/unit_tests/bridge/training/test_setup.py b/tests/unit_tests/bridge/training/test_setup.py
new file mode 100644
--- /dev/null
+++ b/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)
+    assert any("Trainable percentage: 33.33%" in msg for msg in logs)

Contributor guidelines

Per this repo's CONTRIBUTING.md:

• All commits are signed off (Signed-off-by trailer, DCO).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 525f843

[kamran-nvidia]

@andrewwhitecdw Please address the CI failures, I will retrigger CI afterwards

[kamran-nvidia]

/ok to test 188c86f

[kamran-nvidia]

/ok to test ae35199

[kamran-nvidia]

@andrewwhitecdw Please address the CI failures