Does `Muon` save the `state["momentum_buffer"]` correctly in a distributed setting?

[EIFY]

Hi, I still need to test this but if I am reading this correctly

Muon/muon.py

Lines 79 to 94 in f90a42b

	for group in self.param_groups:
	params = group["params"]
	params_pad = params + [torch.empty_like(params[-1])] * (dist.get_world_size() - len(params) % dist.get_world_size())
	for base_i in range(len(params))[::dist.get_world_size()]:
	if base_i + dist.get_rank() < len(params):
	p = params[base_i + dist.get_rank()]
	if p.grad is None:
	# continue
	p.grad = torch.zeros_like(p) # Force synchronization
	state = self.state[p]
	if len(state) == 0:
	state["momentum_buffer"] = torch.zeros_like(p)
	update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"])
	p.mul_(1 - group["lr"] * group["weight_decay"])
	p.add_(update.reshape(p.shape), alpha=-group["lr"])
	dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()], params_pad[base_i + dist.get_rank()])

It seems that each GPU updates state["momentum_buffer"] of its assigned parameters in-place and the parameters themselves are synced with all_gather(), but state["momentum_buffer"] isn't synced. Unless there is some automatic guarantee in place what may happen is that it works fine as long as there is no training interruption, but once you save & load only the state["momentum_buffer"] of the parameters handled by cuda:0 is correct...

[mcleish7]

I agree, this seems to mostly fix my saving and reloading for MuonWithAuxAdam

def gather(self):
    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
        return
    world = torch.distributed.get_world_size()

    for group in self.param_groups:
        if not group["use_muon"]:
            continue

        params = group["params"]
        # Make sure every rank has a buffer tensor to receive the broadcast.
        for p in params:
            st = self.state[p]
            if "momentum_buffer" not in st:
                st["momentum_buffer"] = torch.zeros_like(p)

        # For param i, the canonical copy lives on rank (i % world).
        for i, p in enumerate(params):
            src = i % world
            torch.distributed.broadcast(self.state[p]["momentum_buffer"], src=src)

optimizer.register_state_dict_pre_hook(gather)

Side note: this is chatgpt written based off of @EIFY's branch as the one on there threw a torch error I couldn't understand easily.

On another note, my restarting logic still isn't as reproducible as it is with AdamW alone. Maybe the sorting call in the init is also causing problems? It could be good to add in an option to override the sorting logic so it can be deterministic, decided by the user?

@KellerJordan Have you ran any distributed restarting testing for Muon?

Thanks

[EIFY]

@mcleish7 What torch error did you see? I haven't tried running that dist_test.py yet since I don't have access to a multi-GPU machine at the moment.

[mcleish7]

@mcleish7 What torch error did you see? I haven't tried running that dist_test.py yet since I don't have access to a multi-GPU machine at the moment.

I used my own training code to test but took the gather method and got this error from it.

`ValueError: Argument `gather_list must NOT be specified on non-destination ranks.

[EIFY]

@mcleish7 I blind-fixed that error. Maybe try again?

That said, that branch was only intended to demonstrate the issue definitely. For production fix we may want to also use register_state_dict_post_hook() and register_load_state_dict_post_hook() to remove momentum_buffer of the parameters that the ranks are not responsible for.

I don't think sorting in init presents an issue here. The order of parameters should be preserved already.

[EIFY]

Hi all,

I have debugged and run that dist_test.py on a 8xA100 machine. I can confirm that the current version of Muon does not save state["momentum_buffer"] correctly for parameters not on primary / rank 0 / whatever rank you use to save the state_dict in a checkpoint:

ubuntu@64-181-240-219:~/Downloads/Muon$ torchrun dist_test.py
Without state dict pre-hook to gather:
momentum_buffer for parameter id:
0: torch.Size([128, 128]) tensor

With state dict pre-hook to gather:
momentum_buffer for parameter id:
0: torch.Size([128, 128]) tensor
1: None
2: None
3: None
4: None
5: None
6: None
7: None

With state dict pre-hook to actually gather:
momentum_buffer for parameter id:
0: torch.Size([128, 128]) tensor
1: torch.Size([128, 128]) tensor
2: torch.Size([128, 128]) tensor
3: torch.Size([128, 128]) tensor
4: torch.Size([128, 128]) tensor
5: torch.Size([128, 128]) tensor
6: torch.Size([128, 128]) tensor
7: torch.Size([128, 128]) tensor

Furthermore, apparently gather() does not modify the tensors in gather_list so you have to do the assignment yourself and I suspect it's the same for all_gather()'s tensor_list that Muon uses to sync parameters. Is the distributed Muon tested at all?

Nevermind, I just verified that all parameters are updated on all ranks. Apparently all_gather() does modify tensors in tensor_list in-place, unlike gather()!