Merge pull request #2941 from AI-Hypercomputer:xfgu-rl-fixes

Google-ML-Automation · Google-ML-Automation · commit edb14d9301b0 · 2026-01-14T11:23:46.000-08:00
PiperOrigin-RevId: 856292370
diff --git a/src/MaxText/rl/train_rl.py b/src/MaxText/rl/train_rl.py
@@ -205,8 +205,8 @@ def setup_configs_and_devices(argv: list[str]):
         "dcn_data_parallelism": config.num_samplers_slices,
     }
 
-    trainer_config = pyconfig.initialize_pydantic(config, **trainer_update)
-    sampler_config = pyconfig.initialize_pydantic(config, **sampler_update)
+    trainer_config = pyconfig.initialize_pydantic(argv, **trainer_update)
+    sampler_config = pyconfig.initialize_pydantic(argv, **sampler_update)
 
   else:
     raise ValueError("num_trainer_slices and num_samplers_slices should be both -1 or positive")

Original file line number	Diff line number	Diff line change
`@@ -205,8 +205,8 @@ def setup_configs_and_devices(argv: list[str]):`
`205`	`205`	`"dcn_data_parallelism": config.num_samplers_slices,`
`206`	`206`	`}`
`207`	`207`
`208`		`- trainer_config = pyconfig.initialize_pydantic(config, **trainer_update)`
`209`		`- sampler_config = pyconfig.initialize_pydantic(config, **sampler_update)`
	`208`	`+ trainer_config = pyconfig.initialize_pydantic(argv, **trainer_update)`
	`209`	`+ sampler_config = pyconfig.initialize_pydantic(argv, **sampler_update)`
`210`	`210`
`211`	`211`	`else:`
`212`	`212`	`raise ValueError("num_trainer_slices and num_samplers_slices should be both -1 or positive")`