[feat] max vio in sft

src/prime_rl/trainer/model.py

@@ -16,6 +16,22 @@
 from prime_rl.trainer.config import ActivationCheckpointConfig, ModelConfig
 
 
+def is_tt_moe_model(model: nn.Module) -> bool:
+    return hasattr(model.config, "num_experts") or hasattr(model.config, "n_routed_experts")
+
+
+def get_load_balance_stats(model: nn.Module, reset_stats: bool = True) -> dict[str, torch.FloatTensor]:
+    per_layer_max_vio = []
+    for transformer_block in model.model.layers:
+        tokens_per_expert = transformer_block.mlp.tokens_per_expert
+        balanced_load = tokens_per_expert.mean()
+        max_vio = (tokens_per_expert.max() - balanced_load) / balanced_load
+        per_layer_max_vio.append(max_vio.item())
+        if reset_stats:
+            tokens_per_expert.zero_()
+    return {"max_vio": torch.tensor(per_layer_max_vio)}
+
+
 def get_model(config: ModelConfig) -> nn.Module:
     config_model = AutoConfig.from_pretrained(
         config.name, attn_implementation=config.attn, trust_remote_code=config.trust_remote_code

src/prime_rl/trainer/rl/train.py

@@ -24,6 +24,8 @@
     setup_tokenizer,
     reshard_module,
     setup_model,
+    is_tt_moe_model,
+    get_load_balance_stats,
 )
 from prime_rl.trainer.perf import get_perf_counter
 from prime_rl.trainer.utils import (
@@ -274,15 +276,21 @@ def train(config: RLTrainerConfig):
                 recomputed_logprob_errors[micro_step][loss_mask].detach().to("cpu")
             )
 
+            if is_tt_moe_model(model):
+                load_balance_stats = get_load_balance_stats(model)
+                for k, v in load_balance_stats.items():
+                    tensors[k].append(v)
+
             # Add loss tensors to tensor dict for logging purposes
             for key, loss_tensor in loss_tensors.items():
                 loss_tensor = loss_tensor.detach()[loss_mask].detach().to("cpu")
                 tensors[key].append(loss_tensor)
 
             # Debug log with *local, micro step* stats
-            logger.debug(
-                f"Micro Step {micro_step} | Loss: {tensors['loss'][-1].mean().item():.4f} | Entropy: {tensors['entropy'][-1].mean().item():.4f} | Importance Ratio: {tensors['importance_ratio'][-1].mean().item():.4f}"
-            )
+            micro_step_message = f"Micro Step {micro_step} | Loss: {tensors['loss'][-1].mean().item():.4f} | Entropy: {tensors['entropy'][-1].mean().item():.4f} | Importance Ratio: {tensors['importance_ratio'][-1].mean().item():.4f}"
+            if "max_vio" in tensors:
+                micro_step_message += f" | Max Vio: {tensors['max_vio'][-1].mean().item():.4f}"
+            logger.debug(micro_step_message)
 
         # Optionally, clip the gradients
         logger.debug(f"Clipping gradients to {config.optim.max_norm}")
@@ -329,6 +337,8 @@ def train(config: RLTrainerConfig):
         step_time = time.time() - step_start_time
         current_lr = optimizer.param_groups[0]["lr"]
         step_message = f"Step {progress.step} | Time: {step_time:.2f}s | Loss: {tensor_stats['loss/mean']:.4f} | Entropy: {tensor_stats['entropy/mean']:.4f} | Importance Ratio: {tensor_stats['importance_ratio/mean']:.4f} | Grad. Norm: {grad_norm:.4f} | LR: {current_lr:.2e} | Throughput: {throughput:.0f} tokens/s | MFU: {mfu:.1f}%"
+        if "max_vio/mean" in tensor_stats:
+            step_message += f" | Max Vio: {tensor_stats['max_vio/mean']:.4f}"
         logger.success(step_message)
 
         # Log performance metrics

src/prime_rl/trainer/sft/train.py

@@ -16,6 +16,8 @@
     forward,
     setup_tokenizer,
     setup_model,
+    is_tt_moe_model,
+    get_load_balance_stats,
 )
 from prime_rl.trainer.perf import get_perf_counter
 from prime_rl.trainer.sft.data import setup_dataloader, setup_dataset
@@ -142,6 +144,11 @@ def train(config: SFTTrainerConfig):
             pred_ids = probs.argmax(dim=-1)
             accuracy = torch.eq(pred_ids, target_ids).float()
 
+            if is_tt_moe_model(model):
+                load_balance_stats = get_load_balance_stats(model)
+                for k, v in load_balance_stats.items():
+                    tensors[k].append(v)
+
             # Add tensors to tensor dict for logging purposes
             tensors["loss"].append(loss[loss_mask].detach().to("cpu"))
             tensors["accuracy"].append(accuracy[loss_mask].detach().to("cpu"))
@@ -159,9 +166,10 @@ def train(config: SFTTrainerConfig):
             loss.backward()
 
             # Debug log with *local, micro step* stats
-            logger.debug(
-                f"Micro Step {micro_step} | Loss: {tensors['loss'][-1].mean().item():.4f} | Accuracy: {tensors['accuracy'][-1].mean().item():.4f}"
-            )
+            micro_step_message = f"Micro Step {micro_step} | Loss: {tensors['loss'][-1].mean().item():.4f} | Accuracy: {tensors['accuracy'][-1].mean().item():.4f}"
+            if "max_vio" in tensors:
+                micro_step_message += f" | Max Vio: {tensors['max_vio'][-1].mean().item():.4f}"
+            logger.debug(micro_step_message)
 
         # Optionally, clip the gradients
         logger.debug(f"Clipping gradients to {config.optim.max_norm}")
@@ -201,6 +209,8 @@ def train(config: SFTTrainerConfig):
         step_time = time.time() - step_start_time
         current_lr = optimizer.param_groups[0]["lr"]
         step_message = f"Step {progress.step} | Time: {step_time:.2f}s | Loss: {tensor_stats['loss/mean']:.4f} | Accuracy: {tensor_stats['accuracy/mean']:.4f} | Grad. Norm: {grad_norm:.4f} | LR: {current_lr:.2e} | Throughput: {throughput:.0f} tokens/s | MFU: {mfu:.1f}%"
+        if "max_vio/mean" in tensor_stats:
+            step_message += f" | Max Vio: {tensor_stats['max_vio/mean']:.4f}"
         logger.success(step_message)
 
         # Log progress metrics