update muon_sharding_optimizer with rebuilding 2d_params.

Author: xxyux

test/collective/fleet/hybrid_parallel_sharding_muon_model.py

@@ -243,9 +243,9 @@ def train_batch(self, batch, model, optimizer):
             output = model(batch)
             loss = output.mean()
         loss.backward()
-        if isinstance(optimizer, MuonShardingOptimizer):
+        inner_opt = getattr(optimizer, '_inner_opt', optimizer)
+        if isinstance(inner_opt, MuonShardingOptimizer):
             optimizer.clear_param_storage('test_color')
-            optimizer.reset_param_storage()
         optimizer.step()
         optimizer.clear_grad()
         return loss