just do some shit like that and you should be fine:

  try:
      import torch.distributed as dist
      HAS_DIST = True
  except ImportError:
      HAS_DIST = False
      dist = None

  # In your Muon step() method:
  if HAS_DIST and dist.is_available() and dist.is_initialized():
      # distributed training path
      # ... distributed logic here
  else:
      # single device path - just process parameters normally
      for p in group["params"]:
          if p.grad is None:
              continue
          # apply muon update directly without dist.all_gather