Spaces:
Running on L40S
Running on L40S
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| from __future__ import annotations | |
| import torch | |
| import torch.distributed as dist | |
| import wandb | |
| from cosmos_framework.model._base import ImaginaireModel | |
| from cosmos_framework.utils import distributed | |
| from cosmos_framework.utils.callback import Callback | |
| from cosmos_framework.callbacks.wandb_log import _LossRecord | |
| from cosmos_framework.data.vfm.action.domain_utils import EMBODIMENT_TO_DOMAIN_ID | |
| # Build inverse mapping: domain_id -> embodiment_type. First occurrence wins when multiple embodiment names share the | |
| # same domain id. | |
| DOMAIN_ID_TO_EMBODIMENT: dict[int, str] = {} | |
| for _k, _v in EMBODIMENT_TO_DOMAIN_ID.items(): | |
| DOMAIN_ID_TO_EMBODIMENT.setdefault(_v, _k) | |
| class TrainingStatsCallback(Callback): | |
| """Callback for tracking and logging training mode and embodiment statistics to wandb.""" | |
| def __init__(self, log_freq: int = 100): | |
| super().__init__() | |
| self.log_freq = log_freq | |
| self._mode_counts: dict[str, int] = {} | |
| self._mode_total_count: int = 0 | |
| self._embodiment_counts: dict[str, int] = {} | |
| self._embodiment_total_count: int = 0 | |
| self._per_embodiment_loss: dict[str, _LossRecord] = {} | |
| self._per_embodiment_sub_loss: dict[str, dict[str, _LossRecord]] = {} | |
| def _accumulate_mode_counts(self, data_batch: dict[str, torch.Tensor]) -> None: | |
| modes = data_batch.get("mode", None) | |
| if modes is None: | |
| return | |
| if isinstance(modes, str): | |
| modes_list = [modes] | |
| elif isinstance(modes, (list, tuple)): | |
| modes_list = [str(m) for m in modes] | |
| elif isinstance(modes, torch.Tensor): | |
| # Defensive: support cases where mode might be encoded numerically. | |
| modes_list = [str(m) for m in modes.detach().cpu().tolist()] | |
| else: | |
| modes_list = [str(modes)] | |
| for mode in modes_list: | |
| self._mode_total_count += 1 | |
| self._mode_counts[mode] = self._mode_counts.get(mode, 0) + 1 | |
| def _accumulate_embodiment_counts(self, data_batch: dict[str, torch.Tensor]) -> None: | |
| domain_ids = data_batch.get("domain_id", None) | |
| if domain_ids is None: | |
| return | |
| if isinstance(domain_ids, int): | |
| domain_id_list = [domain_ids] | |
| elif isinstance(domain_ids, (list, tuple)): | |
| domain_id_list = [int(d) for d in domain_ids if d is not None] | |
| elif isinstance(domain_ids, torch.Tensor): | |
| # Flatten to handle any shape (scalar, 1D, or 2D with trailing dim) | |
| domain_id_list = [int(d) for d in domain_ids.detach().cpu().flatten().tolist()] | |
| else: | |
| domain_id_list = [int(domain_ids)] | |
| for domain_id in domain_id_list: | |
| embodiment = DOMAIN_ID_TO_EMBODIMENT.get(domain_id, f"unknown_{domain_id}") | |
| self._embodiment_total_count += 1 | |
| self._embodiment_counts[embodiment] = self._embodiment_counts.get(embodiment, 0) + 1 | |
| def _gather_global_mode_counts(self) -> tuple[int, dict[str, int]]: | |
| """ | |
| Returns (global_total, global_mode_counts) aggregated across all ranks. | |
| """ | |
| local: dict[str, int] = dict(self._mode_counts) | |
| local["__total__"] = int(self._mode_total_count) | |
| if dist.is_available() and dist.is_initialized(): | |
| world_size = int(dist.get_world_size()) | |
| gathered: list[dict[str, int] | None] = [None for _ in range(world_size)] | |
| dist.all_gather_object(gathered, local) | |
| else: | |
| gathered = [local] | |
| global_total = 0 | |
| global_counts: dict[str, int] = {} | |
| for item in gathered: | |
| if not item: | |
| continue | |
| global_total += int(item.get("__total__", 0)) | |
| for k, v in item.items(): | |
| if k == "__total__": | |
| continue | |
| global_counts[k] = global_counts.get(k, 0) + int(v) | |
| return global_total, global_counts | |
| def _gather_global_embodiment_counts(self) -> tuple[int, dict[str, int]]: | |
| """ | |
| Returns (global_total, global_embodiment_counts) aggregated across all ranks. | |
| """ | |
| local: dict[str, int] = dict(self._embodiment_counts) | |
| local["__total__"] = int(self._embodiment_total_count) | |
| if dist.is_available() and dist.is_initialized(): | |
| world_size = int(dist.get_world_size()) | |
| gathered: list[dict[str, int] | None] = [None for _ in range(world_size)] | |
| dist.all_gather_object(gathered, local) | |
| else: | |
| gathered = [local] | |
| global_total = 0 | |
| global_counts: dict[str, int] = {} | |
| for item in gathered: | |
| if not item: | |
| continue | |
| global_total += int(item.get("__total__", 0)) | |
| for k, v in item.items(): | |
| if k == "__total__": | |
| continue | |
| global_counts[k] = global_counts.get(k, 0) + int(v) | |
| return global_total, global_counts | |
| def _build_mode_log_dict( | |
| self, *, log_prefix: str, global_total: int, global_counts: dict[str, int] | |
| ) -> dict[str, float]: | |
| info: dict[str, float] = {} | |
| denom = float(global_total) if global_total > 0 else 0.0 | |
| for mode in sorted(global_counts.keys()): | |
| count = float(global_counts.get(mode, 0)) | |
| pct = (100.0 * count / denom) if denom > 0 else 0.0 | |
| info[f"{log_prefix}_stats_mode/{mode}"] = pct | |
| return info | |
| def _build_embodiment_log_dict( | |
| self, *, log_prefix: str, global_total: int, global_counts: dict[str, int] | |
| ) -> dict[str, float]: | |
| info: dict[str, float] = {} | |
| denom = float(global_total) if global_total > 0 else 0.0 | |
| for embodiment in sorted(global_counts.keys()): | |
| count = float(global_counts.get(embodiment, 0)) | |
| pct = (100.0 * count / denom) if denom > 0 else 0.0 | |
| info[f"{log_prefix}_stats_embodiment/{embodiment}"] = pct | |
| return info | |
| def _get_batch_embodiment(self, data_batch: dict[str, torch.Tensor]) -> str | None: | |
| """Extract the embodiment name from the first non-None sample's domain_id.""" | |
| domain_ids = data_batch.get("domain_id", None) | |
| if domain_ids is None: | |
| return None | |
| if isinstance(domain_ids, torch.Tensor): | |
| if domain_ids.numel() == 0: | |
| return None | |
| domain_id = int(domain_ids.flatten()[0].item()) | |
| elif isinstance(domain_ids, (list, tuple)): | |
| first = next((d for d in domain_ids if d is not None), None) | |
| if first is None: | |
| return None | |
| domain_id = int(first) | |
| else: | |
| domain_id = int(domain_ids) | |
| return DOMAIN_ID_TO_EMBODIMENT.get(domain_id, f"unknown_{domain_id}") | |
| def _accumulate_per_embodiment_loss( | |
| self, | |
| data_batch: dict[str, torch.Tensor], | |
| output_batch: dict[str, torch.Tensor], | |
| loss: torch.Tensor, | |
| ) -> None: | |
| embodiment = self._get_batch_embodiment(data_batch) | |
| if embodiment is None: | |
| return | |
| if embodiment not in self._per_embodiment_loss: | |
| self._per_embodiment_loss[embodiment] = _LossRecord() | |
| self._per_embodiment_loss[embodiment].loss += loss.detach().float() | |
| self._per_embodiment_loss[embodiment].iter_count += 1 | |
| if embodiment not in self._per_embodiment_sub_loss: | |
| self._per_embodiment_sub_loss[embodiment] = {} | |
| for key in output_batch: | |
| if "loss" in key and "per_instance" not in key: | |
| if key not in self._per_embodiment_sub_loss[embodiment]: | |
| self._per_embodiment_sub_loss[embodiment][key] = _LossRecord() | |
| self._per_embodiment_sub_loss[embodiment][key].loss += output_batch[key].detach().float() | |
| self._per_embodiment_sub_loss[embodiment][key].iter_count += 1 | |
| def _compute_per_embodiment_loss_stats(self, log_prefix: str) -> dict[str, float]: | |
| """Compute per-embodiment loss averages across all ranks. | |
| All ranks must call this method (contains collective operations). | |
| Returns the log dict (only meaningful on rank 0). | |
| """ | |
| dist_available = dist.is_available() and dist.is_initialized() | |
| world_size = int(dist.get_world_size()) if dist_available else 1 | |
| # Step 1: gather union of embodiment names across ranks | |
| local_embodiments = sorted(self._per_embodiment_loss.keys()) | |
| if dist_available: | |
| all_embodiments: list[list[str] | None] = [None for _ in range(world_size)] | |
| dist.all_gather_object(all_embodiments, local_embodiments) | |
| else: | |
| all_embodiments = [local_embodiments] | |
| union_embodiments = sorted({e for el in all_embodiments for e in el}) | |
| # Step 2: gather union of sub-loss keys across ranks | |
| local_sub_keys = sorted({k for d in self._per_embodiment_sub_loss.values() for k in d}) | |
| if dist_available: | |
| all_sub_keys: list[list[str] | None] = [None for _ in range(world_size)] | |
| dist.all_gather_object(all_sub_keys, local_sub_keys) | |
| else: | |
| all_sub_keys = [local_sub_keys] | |
| union_sub_keys = sorted({k for kl in all_sub_keys for k in kl}) | |
| # Step 3: insert NaN dummy _LossRecord for missing embodiment/key combos | |
| for emb in union_embodiments: | |
| if emb not in self._per_embodiment_loss: | |
| dummy = _LossRecord() | |
| dummy.loss += torch.tensor([float("nan")], device="cuda") | |
| dummy.iter_count += 1 | |
| self._per_embodiment_loss[emb] = dummy | |
| if emb not in self._per_embodiment_sub_loss: | |
| self._per_embodiment_sub_loss[emb] = {} | |
| for key in union_sub_keys: | |
| if key not in self._per_embodiment_sub_loss[emb]: | |
| dummy = _LossRecord() | |
| dummy.loss += torch.tensor([float("nan")], device="cuda") | |
| dummy.iter_count += 1 | |
| self._per_embodiment_sub_loss[emb][key] = dummy | |
| # Step 4: compute distributed averages (all ranks participate in all_reduce) | |
| log_dict: dict[str, float] = {} | |
| for emb in union_embodiments: | |
| avg, valid = self._per_embodiment_loss[emb].get_stat(return_valid_mask_sum=True) | |
| if valid > 0: | |
| log_dict[f"{log_prefix}_stats_loss/{emb}"] = avg | |
| for emb in union_embodiments: | |
| for key in union_sub_keys: | |
| avg, valid = self._per_embodiment_sub_loss[emb][key].get_stat(return_valid_mask_sum=True) | |
| if valid > 0: | |
| log_dict[f"{log_prefix}_stats_loss_detail/{emb}_{key}"] = avg | |
| # Step 5: reset accumulators | |
| self._per_embodiment_loss = {} | |
| self._per_embodiment_sub_loss = {} | |
| return log_dict | |
| def on_training_step_end( | |
| self, | |
| model: ImaginaireModel, | |
| data_batch: dict[str, torch.Tensor], | |
| output_batch: dict[str, torch.Tensor], | |
| loss: torch.Tensor, | |
| iteration: int = 0, | |
| ) -> None: | |
| self._accumulate_mode_counts(data_batch) | |
| self._accumulate_embodiment_counts(data_batch) | |
| self._accumulate_per_embodiment_loss(data_batch, output_batch, loss) | |
| if iteration % self.log_freq != 0: | |
| return | |
| # All ranks must participate in collective operations below. | |
| mode_total, mode_counts = self._gather_global_mode_counts() | |
| embodiment_total, embodiment_counts = self._gather_global_embodiment_counts() | |
| per_embodiment_loss_dict = self._compute_per_embodiment_loss_stats(log_prefix="train") | |
| if not distributed.is_rank0(): | |
| return | |
| if wandb.run is None: | |
| return | |
| log_dict: dict[str, float] = {} | |
| log_dict.update( | |
| self._build_mode_log_dict(log_prefix="train", global_total=mode_total, global_counts=mode_counts) | |
| ) | |
| log_dict.update( | |
| self._build_embodiment_log_dict( | |
| log_prefix="train", global_total=embodiment_total, global_counts=embodiment_counts | |
| ) | |
| ) | |
| log_dict.update(per_embodiment_loss_dict) | |
| wandb.log({k: float(v) for k, v in log_dict.items()}, step=iteration) | |