Update nemo_rl/models/megatron/setup.py

seonjinn · terrykong · web-flow · commit 55bc3550dba9 · 2026-04-23T14:32:14.000-07:00
Co-authored-by: Terry Kong &lt;terryk@nvidia.com&gt;
Signed-off-by: Seonjin  &lt;sna@nvidia.com&gt;
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
@@ -524,11 +524,12 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None:
     )
     if fine_grained_activation_offloading:
         offload_modules = config["megatron_cfg"].get("offload_modules", [])
-        if not offload_modules:
+        if not isinstance(offload_modules, list) or not offload_modules:
             raise ValueError(
                 "offload_modules must be a non-empty list when "
                 "fine_grained_activation_offloading is True."
             )
+            )
         model_cfg.fine_grained_activation_offloading = True
         model_cfg.offload_modules = offload_modules
 

Original file line number	Diff line number	Diff line change
`@@ -524,11 +524,12 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None:`
`524`	`524`	`)`
`525`	`525`	`if fine_grained_activation_offloading:`
`526`	`526`	`offload_modules = config["megatron_cfg"].get("offload_modules", [])`
`527`		`- if not offload_modules:`
	`527`	`+ if not isinstance(offload_modules, list) or not offload_modules:`
`528`	`528`	`raise ValueError(`
`529`	`529`	`"offload_modules must be a non-empty list when "`
`530`	`530`	`"fine_grained_activation_offloading is True."`
`531`	`531`	`)`
	`532`	`+ )`
`532`	`533`	`model_cfg.fine_grained_activation_offloading = True`
`533`	`534`	`model_cfg.offload_modules = offload_modules`
`534`	`535`