fix: pass layer_wise_distributed_optimizer via config instead of kwarg

The test was passing layer_wise_distributed_optimizer as a keyword arg
to get_megatron_muon_optimizer(), but that function doesn't accept it.
Set it on the OptimizerConfig object instead.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ilml

tests/unit_tests/test_layer_wise_optimizer.py

@@ -124,11 +124,11 @@ def create_model_and_optimizer(
         pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True)
         pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group()
 
+        optimizer_config.use_layer_wise_distributed_optimizer = use_layer_wise
         optimizer = get_megatron_muon_optimizer(
             config=optimizer_config,
             model_chunks=[model],
             use_gloo_process_groups=True,
-            layer_wise_distributed_optimizer=use_layer_wise,
             pg_collection=pg_collection,
         )
         return model, optimizer, pg_collection
@@ -197,11 +197,11 @@ def create_model_and_optimizer_with_overlap_param_gather(
         pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True)
         pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group()
 
+        optimizer_config.use_layer_wise_distributed_optimizer = True
         optimizer = get_megatron_muon_optimizer(
             config=optimizer_config,
             model_chunks=[model],
             use_gloo_process_groups=True,
-            layer_wise_distributed_optimizer=True,
             pg_collection=pg_collection,
         )
         return model, optimizer, pg_collection
@@ -399,11 +399,11 @@ def test_bf16_error(self):
             use_distributed_optimizer=False,
             muon_tp_mode="duplicated",
         )
+        optimizer_config.use_layer_wise_distributed_optimizer = False
         muon_optimizer = get_megatron_muon_optimizer(
             config=optimizer_config,
             model_chunks=[model],
             use_gloo_process_groups=True,
-            layer_wise_distributed_optimizer=False,
             pg_collection=pg_collection,
         )