Add comments and description in the MegatronConfig(TypedDict)

Signed-off-by: Kate Cheng <[email protected]>

Author: katec846

examples/configs/grpo_math_1B.yaml

@@ -85,7 +85,7 @@ policy:
     moe_permute_fusion: false
     #gives ~20% training perf speedup with sequence packing
     apply_rope_fusion: True
-    # gives training perf speedup
+    # gives ~25% training perf speedup with sequence packing and apply_rope_fusion
     bias_activation_fusion: True
     defer_fp32_logits: null
 

examples/configs/sft.yaml

@@ -90,7 +90,7 @@ policy:
     moe_permute_fusion: false
     #gives ~20% training perf speedup with sequence packing 
     apply_rope_fusion: True
-    # gives training perf speedup
+    # gives ~25% training perf speedup with sequence packing and apply_rope_fusion
     bias_activation_fusion: True
 
     optimizer:

examples/configs/sft_openmathinstruct2_megatron.yaml

@@ -85,7 +85,7 @@ policy:
     moe_permute_fusion: false
     #gives ~20% training perf speedup with sequence packing
     apply_rope_fusion: True
-    # gives training perf speedup
+    # gives ~25% training perf speedup with sequence packing and apply_rope_fusion
     bias_activation_fusion: True
 
     env_vars:

nemo_rl/models/policy/__init__.py

@@ -99,6 +99,10 @@ class MegatronConfig(TypedDict):
     expert_tensor_parallel_size: int
     expert_model_parallel_size: int
     defer_fp32_logits: NotRequired[bool]
+    # gives ~20% training perf speedup with sequence packing 
+    apply_rope_fusion: bool
+    # gives ~25% training perf speedup with sequence packing and apply_rope_fusion
+    bias_activation_fusion: bool
 
     optimizer: NotRequired[MegatronOptimizerConfig]
     scheduler: NotRequired[MegatronSchedulerConfig]