Spaces:
Running on L40S
Running on L40S
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| import matplotlib.pyplot as plt | |
| import torch | |
| import wandb | |
| from torch.distributed.tensor import DTensor, Partial | |
| from cosmos_framework.callbacks.every_n import EveryN | |
| from cosmos_framework.model._base import ImaginaireModel | |
| from cosmos_framework.trainer import ImaginaireTrainer | |
| from cosmos_framework.utils import distributed | |
| from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock | |
| def compute_expert_heatmap(vfm: torch.nn.Module) -> dict[str, torch.Tensor]: | |
| """ | |
| Compute the heatmap for the MoE blocks in the language model. | |
| The heatmap is a dictionary with keys set to ["und", "gen"] and values set to | |
| a tensor of shape (num_layers, num_experts). | |
| Each element of the tensor is the average number of tokens routed to each expert for a | |
| given layer. The sum of the elements in each row should be equal to the average number | |
| of experts per token for the MoE model (config.num_experts_per_tok). | |
| For dense models, the heatmap is an empty dictionary. | |
| """ | |
| with torch.no_grad(): | |
| num_layers = len(vfm.language_model.model.layers) | |
| example_dtensor = vfm.language_model.model.layers[0].self_attn.q_proj.weight | |
| if isinstance(example_dtensor, DTensor): | |
| assert hasattr(example_dtensor, "device_mesh") | |
| device_mesh = example_dtensor.device_mesh | |
| else: | |
| device_mesh = None | |
| expert_heatmaps = {} | |
| for tower in ["und", "gen"]: | |
| expert_heatmaps_per_layer = [] | |
| for layer_idx in range(num_layers): | |
| layer_module = vfm.language_model.model.layers[layer_idx] | |
| mlp_module = layer_module.mlp if tower == "und" else layer_module.mlp_moe_gen | |
| if isinstance(mlp_module, Qwen3VLMoeTextSparseMoeBlock): | |
| # This is accumulated across all iterations. | |
| total_tokens_per_expert = mlp_module.get_total_tokens_per_expert() | |
| total_tokens = mlp_module.get_total_tokens() | |
| # Compute the average across all ranks. | |
| assert device_mesh is not None, "MoE models require multiple GPUs." | |
| total_tokens_per_expert = DTensor.from_local( | |
| total_tokens_per_expert, | |
| device_mesh=device_mesh, | |
| placements=[Partial()] * device_mesh.ndim, | |
| ).full_tensor() | |
| total_tokens = DTensor.from_local( | |
| total_tokens, | |
| device_mesh=device_mesh, | |
| placements=[Partial()] * device_mesh.ndim, | |
| ).full_tensor() | |
| mean_tokens_per_expert = total_tokens_per_expert.float() / total_tokens.float() # [num_experts] | |
| expert_heatmaps_per_layer.append(mean_tokens_per_expert) | |
| if len(expert_heatmaps_per_layer) > 0: | |
| expert_heatmaps[tower] = torch.stack(expert_heatmaps_per_layer, dim=0) # [num_layers,num_experts] | |
| return expert_heatmaps | |
| class ExpertHeatmap(EveryN): | |
| """ | |
| Plots the expert heatmap for the MoE blocks in the language model. | |
| Args: | |
| every_n (int): Number of iterations to log the expert heatmap. | |
| """ | |
| def __init__(self, every_n: int = 1000): | |
| super().__init__(every_n=every_n) | |
| def every_n_impl( | |
| self, | |
| trainer: ImaginaireTrainer, | |
| model: ImaginaireModel, | |
| data_batch: dict[str, torch.Tensor], | |
| output_batch: dict[str, torch.Tensor], | |
| loss: torch.Tensor, | |
| iteration: int, | |
| ) -> None: | |
| expert_heatmaps = compute_expert_heatmap(model.net) | |
| if distributed.is_rank0() and wandb.run: | |
| for tower, heatmap in expert_heatmaps.items(): | |
| fig, ax = plt.subplots() | |
| im = ax.imshow(heatmap.cpu().numpy()) | |
| ax.set_xlabel("Experts") | |
| ax.set_ylabel("Layers") | |
| plt.colorbar(im, ax=ax) | |
| wandb.log( | |
| { | |
| f"expert_heatmap/{tower}": fig, | |
| }, | |
| step=iteration, | |
| ) | |
| plt.close(fig) | |