cp: [dev] Fix cuda graph scope check in language_model.py (NVIDIA#2158) (NVIDIA#2159)

ananthsub · web-flow · commit c69e6c38c64e · 2025-11-06T21:53:05.000+01:00
Signed-off-by: Ananth Subramaniam &lt;ansubramania@nvidia.com&gt;
diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
@@ -138,6 +138,7 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor:
                     # Use is_cg_capturable=True for full iteration CUDA graphs to avoid torch.equal checks
                     is_cg_capturable = (
                         hasattr(self.config, 'cuda_graph_scope')
+                        and self.config.cuda_graph_scope
                         and 'full_iteration' in self.config.cuda_graph_scope
                     )
                     if is_cg_capturable and not is_te_min_version("2.7.0"):

Original file line number	Diff line number	Diff line change
`@@ -138,6 +138,7 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor:`
`138`	`138`	`# Use is_cg_capturable=True for full iteration CUDA graphs to avoid torch.equal checks`
`139`	`139`	`is_cg_capturable = (`
`140`	`140`	`hasattr(self.config, 'cuda_graph_scope')`
	`141`	`+ and self.config.cuda_graph_scope`
`141`	`142`	`and 'full_iteration' in self.config.cuda_graph_scope`
`142`	`143`	`)`
`143`	`144`	`if is_cg_capturable and not is_te_min_version("2.7.0"):`