fix(megatron): patch validate_args and add ROCM argument validation

Hook Megatron validate_args alongside parse_args so Primus-injected arguments are validated consistently, and run additional ROCM-specific argument checks during initialization.

Author: WangLingxun

primus/backends/megatron/megatron_base_trainer.py

@@ -23,17 +23,35 @@ def init(self):
         # log_dict_aligned("Backend arguments", self.backend_args)
 
     def _patch_parse_args(self):
-        """Patch Megatron's parse_args to return pre-configured Primus arguments."""
+        """
+        This function patches Megatron's parse_args to return pre-configured Primus arguments.
+        It also validates the arguments on ROCM.
+        """
         import megatron.training.arguments as megatron_args  # type: ignore
         import megatron.training.initialize as megatron_init  # type: ignore
+        from primus.modules.trainer.megatron.utils import validate_args_on_rocm
 
         log_rank_0("Patching Megatron-LM parse_args()")
 
         patched_parse_args = lambda *args, **kwargs: (
             log_rank_0("parse_args() called; returning Primus arguments") or self.backend_args
         )
 
+        original_validate_args = megatron_args.validate_args
+        def patched_validate_args(*args, **kwargs):
+            validated_args = original_validate_args(*args, **kwargs)
+            parsed_args = args[0] if args else kwargs.get("args", None)
+            if parsed_args is not None:
+                log_rank_0("validate_args() called; validating on ROCM")
+                validate_args_on_rocm(parsed_args)
+            return validated_args
+
         megatron_args.parse_args = patched_parse_args
         megatron_init.parse_args = patched_parse_args
 
-        log_rank_0(f"Patched parse_args(); Primus provided {len(vars(self.backend_args))} arguments")
+        megatron_args.validate_args = patched_validate_args
+        megatron_init.validate_args = patched_validate_args
+
+        log_rank_0(
+            f"Patched parse_args()/validate_args(); Primus provided {len(vars(self.backend_args))} arguments"
+        )