diff --git a/docs/usage/trainer.md b/docs/usage/trainer.md index 4f10a7814b..88d14a458d 100644 --- a/docs/usage/trainer.md +++ b/docs/usage/trainer.md @@ -213,3 +213,19 @@ class Arguments(VeOmniArguments): ``` By following this pattern, you can leverage the robust infrastructure of `BaseTrainer` while tailoring the training process to your specific model and data requirements. + +### Token accounting diagnostics + +When the batch contains sufficient global length and padding metadata, +`EnvironMeter` reports `token_accounting/*` diagnostics. The physical count is +the SP-padded input window, the aligned count removes its known tail padding, +and the source count uses valid global sequence lengths. Ambiguous batches omit +these metrics. Existing `consume_tokens` and multi-source dataloader accounting +retain their previous behavior and can therefore differ from the new source +count. + +`loss_tokens` is a generic shifted-label count, not an objective-specific loss +denominator. In particular, do not compare it across SP configurations for +objectives whose collator does not shift labels, or treat it as the exact DPO, +classification, or MTP target count. Token throughput uses the step timer, +which starts after meter input collection and does not include its overhead. diff --git a/tests/data/test_token_meter.py b/tests/data/test_token_meter.py new file mode 100644 index 0000000000..429f7339d7 --- /dev/null +++ b/tests/data/test_token_meter.py @@ -0,0 +1,237 @@ +"""Topology and shifted-label accounting without changing model inputs.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.utils import helper + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +def test_packed_lengths_keep_real_single_token_documents(): + batch = { + "input_ids": torch.arange(4).unsqueeze(0), + "labels": torch.arange(4).unsqueeze(0), + "cu_seq_lens_q": torch.tensor([0, 3, 4]), + "linear_attn_cu_seq_lens_q": torch.tensor([0, 3, 4]), + "packed_sequence_slices": [(0, 3), (3, 4)], + } + meter = _meter(SimpleNamespace(sp_enabled=False, dp_group=None)) + meter._record_token_stats(batch) + assert meter._batch_token_stats[0][3] == 4 + # Existing dataloader accounting keeps its prior behavior. + assert helper._compute_seqlens(batch) == [3] + batch["tail_padding_length"] = torch.tensor(1) + meter = _meter(SimpleNamespace(sp_enabled=False, dp_group=None)) + meter._record_token_stats(batch) + assert meter._batch_token_stats[0][3] == 3 + + +def _meter(state): + meter = object.__new__(helper.EnvironMeter) + meter.config = SimpleNamespace(condition_model_type=None) + meter.parallel_state = state + meter.global_batch_size = 2 + meter.world_size = 4 + meter.consume_tokens = meter.consume_chunks = 0 + meter.batch_seqlens = [] + meter.batch_ds_idx = [] + meter.images_seqlens = [] + meter._batch_token_stats = [] + meter._batch_token_stats_complete = True + meter.enable_multisource = False + meter.empty_cache_steps = meter.gc_steps = 0 + meter.supports_lora_flops = False + meter.estimate_flops = lambda lengths, delta_time, **kwargs: (0.0, 1.0) + return meter + + +def _sp_worker(rank, rendezvous, missing): + dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=4) + try: + # Two data-parallel replicas, each with two distinct sequence shards. + sp_groups = [dist.new_group(ranks) for ranks in ([0, 1], [2, 3])] + dp_groups = [dist.new_group(ranks) for ranks in ([0, 2], [1, 3])] + state = SimpleNamespace( + sp_enabled=True, sp_size=2, sp_group=sp_groups[rank // 2], dp_group=dp_groups[rank % 2] + ) + meter = _meter(state) + full_labels = torch.tensor([[-100, 1, 2, -100, 3, -100, -100, -100]]) + batch = { + "input_ids": torch.arange(4).view(1, -1), + "labels": full_labels[:, (rank % 2) * 4 : (rank % 2 + 1) * 4], + "cu_seq_lens_q": torch.tensor([0, 2, 4]), + "linear_attn_cu_seq_lens_q": torch.tensor([0, 3, 5, 8]), + "packed_sequence_slices": [(0, 3), (3, 5), (5, 8)], + "tail_padding_length": torch.tensor(3), + } + if missing and rank == 0: + batch.pop("labels") + meter.add(batch) + assert meter.batch_seqlens == [] # Existing local-cu tracker behavior is unchanged. + with ( + patch.object(helper, "compute_device_memory_metrics", return_value={}), + patch.object(helper, "get_device_type", return_value="cpu"), + ): + metrics = meter.step(delta_time=2.0, global_step=1) + if missing: + assert not any(key.startswith("token_accounting/") for key in metrics) + return + assert metrics["token_accounting/physical_window_tokens"] == 16 + assert metrics["token_accounting/aligned_compute_tokens"] == 10 + assert metrics["token_accounting/source_input_tokens"] == 10 + assert metrics["token_accounting/loss_tokens"] == 6 + assert metrics["token_accounting/source_fill"] == pytest.approx(10 / 16) + assert meter._batch_token_stats == [] + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) +@pytest.mark.parametrize("missing", [False, True]) +def test_token_meter_packed_dp2_sp2(tmp_path, missing): + mp.start_processes(_sp_worker, args=(str(tmp_path / "rdzv"), missing), nprocs=4, join=True, start_method="fork") + + +def _real_collator_sp_worker(rank, rendezvous): + from veomni.data import data_collator + + dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=4) + try: + dp_groups = [dist.new_group([member]) for member in range(4)] + state = SimpleNamespace( + sp_enabled=True, sp_size=4, sp_rank=rank, sp_group=dist.group.WORLD, dp_group=dp_groups[rank] + ) + samples = [ + {"input_ids": torch.arange(length), "labels": torch.arange(length), "attention_mask": torch.ones(length)} + for length in (5, 7) + ] + with patch.object(data_collator, "get_parallel_state", return_value=state): + batch = data_collator.MainCollator(pad_to_length=16)(samples) + meter = _meter(state) + meter.add(batch) + assert meter.batch_seqlens == [5, 7] + with ( + patch.object(helper, "compute_device_memory_metrics", return_value={}), + patch.object(helper, "get_device_type", return_value="cpu"), + ): + metrics = meter.step(delta_time=2.0, global_step=1) + assert metrics["token_accounting/physical_window_tokens"] == 16 + assert metrics["token_accounting/aligned_compute_tokens"] == 12 + assert metrics["token_accounting/source_input_tokens"] == 12 + assert metrics["token_accounting/loss_tokens"] == 10 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) +def test_real_collator_packed_sp4(tmp_path): + mp.start_processes( + _real_collator_sp_worker, args=(str(tmp_path / "rdzv-sp4"),), nprocs=4, join=True, start_method="fork" + ) + + +def _actual_mesh_worker(rank, rendezvous, use_cp): + from veomni.data import data_collator + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state + + dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=4) + try: + state = _init_parallel_state( + dp_size=2, + dp_shard_size=2, + cp_size=2 if use_cp else 1, + ulysses_size=1 if use_cp else 2, + device_type="cpu", + ) + lengths = (3, 1) if state.dp_rank == 0 else (5, 2) + samples = [ + { + "input_ids": torch.arange(n), + "labels": torch.arange(n), + "attention_mask": torch.ones(n, dtype=torch.long), + } + for n in lengths + ] + batch = data_collator.MainCollator(pad_to_length=8)(samples) + meter = _meter(state) + meter.add(batch) + with ( + patch.object(helper, "compute_device_memory_metrics", return_value={}), + patch.object(helper, "get_device_type", return_value="cpu"), + ): + metrics = meter.step(delta_time=2.0, global_step=1) + names = ("physical_window_tokens", "aligned_compute_tokens", "source_input_tokens", "loss_tokens") + assert tuple(metrics["token_accounting/" + name] for name in names) == (16, 11, 11, 7) + finally: + dist.destroy_process_group() + clear_parallel_state() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) +@pytest.mark.parametrize("use_cp", [False, True]) +def test_actual_dp_sequence_mesh(tmp_path, use_cp): + mp.start_processes( + _actual_mesh_worker, args=(str(tmp_path / "rdzv-mesh"), use_cp), nprocs=4, join=True, start_method="fork" + ) + + +def test_unsharded_generic_shifted_label_count(): + state = SimpleNamespace(sp_enabled=False, dp_group=None) + meter = _meter(state) + meter._record_token_stats( + { + "input_ids": torch.arange(4).unsqueeze(0), + "labels": torch.tensor([[4, 5, -100, 7]]), + "attention_mask": torch.ones(1, 4), + } + ) + assert meter._batch_token_stats == [(4, 4, 2, 4)] + + +def test_ambiguous_input_is_not_reported_as_exact(): + meter = _meter(SimpleNamespace(sp_enabled=True, sp_size=2, dp_group=None)) + meter._record_token_stats({"input_ids": torch.arange(4), "labels": torch.arange(4)}) + assert not meter._batch_token_stats_complete + assert meter._batch_token_stats == [] + + meter = _meter(SimpleNamespace(sp_enabled=False, dp_group=None)) + meter._record_token_stats({"input_ids": torch.arange(4).unsqueeze(0), "labels": torch.tensor([1])}) + assert not meter._batch_token_stats_complete + assert meter._batch_token_stats == [] + + for metadata in ({}, {"cu_seq_lens_q": torch.tensor([0, 3, 4])}): + meter = _meter(SimpleNamespace(sp_enabled=False, dp_group=None)) + meter._record_token_stats( + {"input_ids": torch.arange(4).unsqueeze(0), "labels": torch.arange(4).unsqueeze(0), **metadata} + ) + assert not meter._batch_token_stats_complete + + meter = _meter(SimpleNamespace(sp_enabled=True, sp_size=2, dp_group=None)) + meter._record_token_stats( + { + "input_ids": torch.arange(4).unsqueeze(0), + "labels": torch.arange(4).unsqueeze(0), + "linear_attn_cu_seq_lens_q": torch.tensor([0, 3, 4]), + "packed_sequence_slices": [(0, 3), (3, 4)], + } + ) + assert not meter._batch_token_stats_complete + + meter = _meter(SimpleNamespace(sp_enabled=False, dp_group=None)) + meter._record_token_stats( + { + "input_ids": torch.arange(8).unsqueeze(0), + "labels": torch.arange(8).unsqueeze(0), + "cu_seq_lens_q": torch.tensor([0, 5, 8]), + "tail_padding_length": torch.tensor(2), + } + ) + assert not meter._batch_token_stats_complete diff --git a/veomni/utils/helper.py b/veomni/utils/helper.py index 6f649efa4d..39064f66ff 100644 --- a/veomni/utils/helper.py +++ b/veomni/utils/helper.py @@ -38,6 +38,7 @@ from ..distributed.parallel_state import get_parallel_state from . import logging +from .constants import IGNORE_INDEX from .count_flops import VeomniFlopsCounter from .device import ( IS_CUDA_AVAILABLE, @@ -246,6 +247,8 @@ def __init__( self.batch_seqlens = [] self.batch_ds_idx = [] self.images_seqlens = [] + self._batch_token_stats = [] + self._batch_token_stats_complete = True if self.enable_multisource: if dataloader is None or data_path is None: @@ -289,16 +292,79 @@ def add(self, micro_batch: Union[Dict[str, "torch.Tensor"], List[Dict[str, "torc for sample in micro_batch: self.batch_seqlens.extend(_compute_seqlens(sample)) self.images_seqlens.extend(_compute_image_seqlens(sample, self.config)) + self._record_token_stats(sample) if self.enable_multisource: self.batch_ds_idx.extend(_get_multisource_ds_idx(sample)) else: self.batch_seqlens.extend(_compute_seqlens(micro_batch)) self.images_seqlens.extend(_compute_image_seqlens(micro_batch, self.config)) + self._record_token_stats(micro_batch) if self.enable_multisource: self.batch_ds_idx.extend(_get_multisource_ds_idx(micro_batch)) else: # dit diffusers model self.batch_seqlens.extend(_compute_wan_seqlens(micro_batch)) + def _record_token_stats(self, batch): + """Meter standard packed batches without changing existing data-cursor accounting. + + Physical counts the padded SP window, aligned removes the known tail, + and source counts valid global sequence lengths. ``loss_tokens`` is + the generic shifted-label count; it is not an objective-specific loss + denominator and is not comparable across SP for unshifted objectives. + Ambiguous metadata is omitted from the new metrics. + """ + ids, labels = batch.get("input_ids"), batch.get("labels") + if ( + not isinstance(ids, torch.Tensor) + or not isinstance(labels, torch.Tensor) + or ids.ndim == 0 + or labels.ndim == 0 + or ids.shape != labels.shape + ): + self._batch_token_stats_complete = False + return + sp_enabled = getattr(self.parallel_state, "sp_enabled", False) + if sp_enabled and "linear_attn_cu_seq_lens_q" not in batch: + # A rank-local boundary list cannot prove a global token count. + self._batch_token_stats_complete = False + return + sp_size = self.parallel_state.sp_size if sp_enabled else 1 + capacity = ids.numel() * sp_size + has_boundaries = "linear_attn_cu_seq_lens_q" in batch or "cu_seq_lens_q" in batch + if has_boundaries: + cu = batch.get("linear_attn_cu_seq_lens_q", batch.get("cu_seq_lens_q")) + trusted_lengths = ( + isinstance(cu, torch.Tensor) + and cu.ndim == 1 + and cu.numel() >= 2 + and int(cu[0]) == 0 + and int(cu[-1]) == capacity + and ("packed_sequence_slices" in batch or batch.get("tail_padding_length") is not None) + ) + else: + mask = batch.get("attention_mask") + trusted_lengths = not sp_enabled and isinstance(mask, torch.Tensor) and mask.shape == ids.shape + if not trusted_lengths: + self._batch_token_stats_complete = False + return + tail = int(batch.get("tail_padding_length", 0) or 0) + if tail < 0 or tail > capacity: + self._batch_token_stats_complete = False + return + if has_boundaries and tail and not bool((cu == capacity - tail).any().item()): + self._batch_token_stats_complete = False + return + source = ( + int(valid_seqlens_from_cu_seqlens(cu, tail_padding_length=tail).sum().item()) + if has_boundaries + else int(mask.sum().item()) + ) + # SP collators already shift labels before slicing. This is a generic + # shifted-label count, not an objective-specific loss denominator. + shifted = labels if sp_enabled else labels[..., 1:] + loss_tokens = int((shifted != IGNORE_INDEX).sum().item()) + self._batch_token_stats.append((capacity, capacity - tail, loss_tokens, source)) + def step( self, delta_time: float, @@ -355,6 +421,49 @@ def step( "consumed_chunk_num": self.consume_chunks, } metrics.update(compute_device_memory_metrics()) + stats = self._batch_token_stats + local_loss = sum(item[2] for item in stats) + complete = int(bool(stats) and self._batch_token_stats_complete) + sp_enabled = getattr(self.parallel_state, "sp_enabled", False) + sp_size = self.parallel_state.sp_size if sp_enabled else 1 + if sp_enabled: + local_loss, complete = all_reduce( + torch.tensor([local_loss, complete], dtype=torch.int64, device=get_device_type()), + op="sum", + group=self.parallel_state.sp_group, + ) + capacity, aligned, loss_tokens, complete_ranks, source_tokens = all_reduce( + torch.tensor( + [ + sum(item[0] for item in stats), + sum(item[1] for item in stats), + local_loss, + complete, + sum(item[3] for item in stats), + ], + dtype=torch.int64, + device=get_device_type(), + ), + op="sum", + group=self.parallel_state.dp_group, + ) + if complete_ranks and complete_ranks == sp_size * dist.get_world_size(group=self.parallel_state.dp_group): + metrics.update( + { + "token_accounting/physical_window_tokens": capacity, + "token_accounting/aligned_compute_tokens": aligned, + "token_accounting/source_input_tokens": source_tokens, + "token_accounting/loss_tokens": loss_tokens, + "token_accounting/source_fill": source_tokens / capacity if capacity else 0.0, + "token_accounting/aligned_compute_fill": aligned / capacity if capacity else 0.0, + "token_accounting/loss_density": loss_tokens / source_tokens if source_tokens else 0.0, + "token_accounting/loss_fill": loss_tokens / capacity if capacity else 0.0, + "token_accounting/capacity_tokens_per_second(M)": capacity / delta_time / 1e6, + "token_accounting/aligned_compute_tokens_per_second(M)": aligned / delta_time / 1e6, + "token_accounting/source_tokens_per_second(M)": source_tokens / delta_time / 1e6, + "token_accounting/loss_tokens_per_second(M)": loss_tokens / delta_time / 1e6, + } + ) if self.enable_multisource: metrics.update(self.multisource_tracker.step(self.batch_ds_idx, self.batch_seqlens)) @@ -368,6 +477,8 @@ def step( self.batch_seqlens = [] self.batch_ds_idx = [] self.images_seqlens = [] + self._batch_token_stats = [] + self._batch_token_stats_complete = True return metrics