[VLM] Add token-imbalance loss

[lkhphuc]

In VLM interleaved training, with native resolution and aspect ratio, the number of tokens participating in loss computation differ per rank. Naive FSDP gradient averaging across data ranks can causes tokens on ranks with fewer valid tokens to contribute more to the loss than on other ranks.
This PR address this via loss balancing, which incur an additional comm in the loss computation.
In practice, I haven't notice any impacts from this comm.

Quick sanity check

Let have a sum loss of all tokens on each rank i, with $N_i$ number of tokens $L_i = \sum_{j=1}^{N_i}\ell_{ij}$ and its gradient $g_i = \sum_{j=1}^{N_i}\nabla\ell_{ij}$

If we multiply the loss on each rank by a constant factor c (the same for all ranks), then after backward():

$$ \tilde g_i = c \cdot g_i . $$

FSDP will average these gradients across ranks:

$$ g_{\text{FSDP}}=\frac{1}{R}\sum_{i=1}^{R} \tilde g_i =\frac{c}{R}\sum_{i=1}^{R} g_i . $$

We want this to equal the global‑sample average:

$$ g_{\text{true}} =\frac{1}{N_{\text{total}}}\sum_{i=1}^{R}\sum_{j=1}^{N_i}\nabla \ell_{ij} =\frac{1}{N_{\text{total}}}\sum_{i=1}^{R} g_i . $$

Thus for FSDP gradient to be correct, we need

$$ \frac{c}{R}= \frac{1}{N_{\text{total}}}\quad\Longrightarrow\quad c=\frac{R}{N_{\text{total}}}. $$

So the right scaling factor is $R/N_{\text{total}}$, which mean divide the per-rank sum loss with $N_{\text{total}}/R$, which is average number of tokens per rank.
Intuitively, this is the same as default cross-entropy loss, but instead of diving sum loss on a rank by the number of tokens on that rank, we now divide by the average number of tokens across all rank

P/s: sorry this PR is based on #1802 but I couldn't choose that as the base branch. Maybe it will be easier to review once that PR is merged.

[tianyu-l]

For this dist_mean call, it seems it'll trigger a GPU/CPU sync

torchtitan/torchtitan/distributed/utils.py

Line 50 in 21739fd

return funcol.all_reduce(x, reduceOp=reduceOp, group=mesh).item()

I think this will potentially bring unnecessary perf issues, as without with CPU can stay way ahead of GPU.

I'd recommend refactoring the .item() call outside _dist_reduce and put them into callsites. Alternatively you can directly call funcol.all_reduce() here.

cc @fegin this won't work with FT as the ft_pg is not visible to this loss function.

[lkhphuc]

I added an extra process group mimicking the distributed utils. Hope that makes sense.

[tianyu-l]

This is a cute & "mostly correct" way to deal with the imbalanced token loss issue. However,

1. it's not the most readable one
2. moreover, I don't think it is correct if gradient accumulation is enabled, as each microbatch can have different amount of "avg_num_tokens"

I think the best way should be

• don't let FSDP do implicit gradient division
• always run cross entropy with reduction="sum"
• let data loader / trainer count the number of tokens involving in loss computation, e.g. by explicitly doing num_tokens = (labels != IGNORE_INDEX).sum() on each rank. (I agree that without imbalance we don't need to do this and the followed communication.)

This way we also don't need this ad hoc call https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/llama4/infra/parallelize.py#L363

We don't need to do this refactor for now, but it would be good if you could leave a TODO item here + file an issue.

cc @ezyang @fegin

[lkhphuc]

I filed #1842

[tianyu-l]

Let's avoid defining IGNORE_INDEX in two places when they actually need to be shared.
The other appearance is in torchtitan/experiments/vlm/datasets/mm_collator_nld.py

[lkhphuc]

I keep the def in loss.py as it makes most sense and import to the dataloader files.

[lkhphuc]

Actually I delete the other PR #1802 from this PR as they are fairly independent to make it cleaner to merge.
I will fix the other PR to import IGNORE_INDEX depend on which ones land first.

[tianyu-l]

Thanks for writing up the docstring, looks very good.

[tianyu-l]

LGTM