feat(optim): optional bfloat16 momentum for Muon

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -439,6 +439,9 @@ class MuonConfig(BaseOptimizerConfig):
     betas2: float = Field(0.95, ge=0)
     """β2 for the AdamW/Lion sub-optimizer used on non-Muon params."""
 
+    bf16_momentum: bool = False
+    """Keep the Muon momentum in bfloat16 instead of the parameter dtype. Halves Muon's optimizer state and its CPU offload traffic, at the cost of momentum precision. Muon already casts the momentum to bfloat16 before Newton-Schulz."""
+
 
 class SignSGDConfig(BaseOptimizerConfig):
     type: Literal["sign_sgd"] = "sign_sgd"

src/prime_rl/trainer/optim/__init__.py

@@ -47,6 +47,24 @@ def _warmup_muon_mesh(mesh: DeviceMesh) -> None:
     get_logger().info(f"Finished warming Muon all-to-all connections on group {group.group_name}")
 
 
+class BF16MomentumMuon(Muon):
+    """Muon whose momentum is bfloat16, also after a checkpoint load (which casts states to the param dtype)."""
+
+    def _get_or_initialize_state(self, param: torch.Tensor, algo: str) -> dict:
+        state = self.state[param]
+        if not state and algo == "muon":
+            state["momentum"] = torch.zeros_like(param, dtype=torch.bfloat16)
+        return super()._get_or_initialize_state(param, algo)
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        super().load_state_dict(state_dict)
+        for group in self.param_groups:
+            if group["algorithm"] == "muon":
+                for param in group["params"]:
+                    if "momentum" in self.state[param]:
+                        self.state[param]["momentum"] = self.state[param]["momentum"].bfloat16()
+
+
 def setup_optimizer(
     config: OptimizerConfig,
     named_params: list[tuple[str, nn.Parameter]],
@@ -219,7 +237,8 @@ def muon_enabled(n, p):
             if partitions is not None and packed_info.parameter in muon_params:
                 matrix_partitions[packed_info.parameter] = partitions
 
-    optimizer = Muon(
+    muon_cls = BF16MomentumMuon if config.bf16_momentum else Muon
+    optimizer = muon_cls(
         params=param_groups,
         matrix_partitions=matrix_partitions,
         lr=lr,