fix: make layernorm_epsilon configurable in with megatron backend (#1046)

ashors1 · web-flow · commit f182460208b0 · 2025-09-03T01:58:15.000Z
Signed-off-by: ashors1 &lt;ashors@nvidia.com&gt;
diff --git a/nemo_rl/models/policy/megatron_policy_worker.py b/nemo_rl/models/policy/megatron_policy_worker.py
@@ -561,6 +561,9 @@ def __init__(
             "moe_router_bias_update_rate"
         ]
 
+        if "layernorm_epsilon" in self.cfg["megatron_cfg"]:
+            model_cfg.layernorm_epsilon = self.cfg["megatron_cfg"]["layernorm_epsilon"]
+
         model_cfg.sequence_parallel = self.cfg["megatron_cfg"]["sequence_parallel"]
         model_cfg.bf16 = self.dtype == torch.bfloat16
         model_cfg.fp16 = self.dtype == torch.float16

Original file line number	Diff line number	Diff line change
`@@ -561,6 +561,9 @@ def __init__(`
`561`	`561`	`"moe_router_bias_update_rate"`
`562`	`562`	`]`
`563`	`563`
	`564`	`+ if "layernorm_epsilon" in self.cfg["megatron_cfg"]:`
	`565`	`+ model_cfg.layernorm_epsilon = self.cfg["megatron_cfg"]["layernorm_epsilon"]`
	`566`	`+`
`564`	`567`	`model_cfg.sequence_parallel = self.cfg["megatron_cfg"]["sequence_parallel"]`
`565`	`568`	`model_cfg.bf16 = self.dtype == torch.bfloat16`
`566`	`569`	`model_cfg.fp16 = self.dtype == torch.float16`