Mixtral 8×7B / 8×22B FT: Router Collapse Problem + Aux Loss Weight Calibration

# === Expert Balance Metrics ölçümü ===
# Mixtral router'ı her step için expert distribution'unu logla
import torch
 
class RouterMonitor:
    def __init__(self, num_experts=8):
        self.num_experts = num_experts
        self.token_counts = torch.zeros(num_experts)
        self.weight_sums = torch.zeros(num_experts)
 
    def update(self, router_logits, expert_indices):
        """Her FFN forward call'unda çağır."""
        with torch.no_grad():
            for i in range(self.num_experts):
                mask = (expert_indices == i)
                self.token_counts[i] += mask.sum().item()
                self.weight_sums[i] += router_logits[mask].sum().item()
 
    def report(self):
        total = self.token_counts.sum()
        if total == 0: return {}
        f = self.token_counts / total
        return {
            "expert_load_ratio": f.tolist(),
            "expert_load_std": f.std().item(),       # 0 = perfect balance
            "expert_load_max": f.max().item(),       # >0.4 → router collapse riski
            "expert_load_min": f.min().item(),       # <0.05 → dead expert
        }
 
# Training callback ile entegre et
from transformers import TrainerCallback
 
class MoEMonitorCallback(TrainerCallback):
    def __init__(self, model):
        self.monitors = {}
        for layer_idx, layer in enumerate(model.model.layers):
            if hasattr(layer, "block_sparse_moe"):
                self.monitors[layer_idx] = RouterMonitor()
 
    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step % 50 == 0:
            for idx, mon in self.monitors.items():
                report = mon.report()
                print(f"layer-{idx}: load_std={report['expert_load_std']:.3f}, max={report['expert_load_max']:.3f}")
 
# Sağlıklı bir Mixtral FT'inde her layer'da:
# - load_std ~0.03-0.08 (uniform'a yakın)
# - load_max ~0.15-0.20 (ideal: 1/8 = 0.125)
# - load_min ~0.08-0.12
 
# Router collapse'da:
# - load_std > 0.15
# - load_max > 0.40 (bir expert tüm token'ları yutuyor)
# - load_min < 0.02 (dead expert)