Fix W&B step mismatch by consolidating log calls (#15204)

The W&B _step counter was showing approximately 2x the actual
global_step because `log('global_step')` and `log('step')` were
called separately, causing WandbLogger's internal _step counter
to increment twice per training step.

This fix consolidates both metrics into a single `log_dict()` call
in all three strategy files, ensuring W&B step aligns correctly
with trainer.global_step for accurate cross-run comparisons.

Files modified:
- nemo/lightning/pytorch/strategies/megatron_strategy.py
- nemo/lightning/pytorch/strategies/fsdp_strategy.py
- nemo/lightning/pytorch/strategies/fsdp2_strategy.py

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: yurekami

nemo/lightning/pytorch/strategies/fsdp2_strategy.py

@@ -401,19 +401,16 @@ def training_step(self, batch, batch_idx=None) -> STEP_OUTPUT:
         else:
             loss = self.lightning_module.training_step(batch, batch_idx)
 
-        self.lightning_module.log(
-            'global_step',
-            self.trainer.global_step,
+        # Log global_step and step together using log_dict to prevent W&B from
+        # incrementing its internal _step counter multiple times per training step.
+        # This fixes issue #15204 where W&B showed ~2x the actual global_step.
+        self.lightning_module.log_dict(
+            {'global_step': self.trainer.global_step, 'step': self.trainer.global_step},
             prog_bar=True,
             rank_zero_only=True,
             batch_size=1,
         )
 
-        self.lightning_module.log(
-            'step',
-            self.trainer.global_step,
-        )
-
         return loss
 
     @override

nemo/lightning/pytorch/strategies/fsdp_strategy.py

@@ -181,18 +181,15 @@ def training_step(self, batch, batch_idx=None) -> STEP_OUTPUT:
         with self.precision_plugin.train_step_context():
             loss, reduced = self._step_proxy("training", batch, batch_idx)
 
-            self.lightning_module.log(
-                'global_step',
-                self.trainer.global_step,
+            # Log global_step and step together using log_dict to prevent W&B from
+            # incrementing its internal _step counter multiple times per training step.
+            # This fixes issue #15204 where W&B showed ~2x the actual global_step.
+            self.lightning_module.log_dict(
+                {'global_step': self.trainer.global_step, 'step': self.trainer.global_step},
                 prog_bar=True,
                 rank_zero_only=True,
                 batch_size=1,
             )
-
-            self.lightning_module.log(
-                'step',
-                self.trainer.global_step,
-            )
             self.lightning_module.log(
                 'reduced_train_loss', reduced['avg'], prog_bar=True, rank_zero_only=True, batch_size=1
             )

nemo/lightning/pytorch/strategies/megatron_strategy.py

@@ -800,18 +800,15 @@ def training_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OUTP
                 if (torch.is_tensor(self.trainer.global_step) and self.trainer.global_step.is_cuda)
                 else torch.tensor(self.trainer.global_step, pin_memory=True).to("cuda", non_blocking=True)
             )
-            self.lightning_module.log(
-                "global_step",
-                global_step,
+            # Log global_step and step together using log_dict to prevent W&B from
+            # incrementing its internal _step counter multiple times per training step.
+            # This fixes issue #15204 where W&B showed ~2x the actual global_step.
+            self.lightning_module.log_dict(
+                {"global_step": global_step, "step": global_step},
                 prog_bar=True,
                 batch_size=1,
             )
 
-            self.lightning_module.log(
-                "step",
-                global_step,
-            )
-
             if self.log_memory_usage:
                 # maximum GPU memory that has been managed by the caching allocator
                 max_memory_reserved = torch.cuda.max_memory_reserved()