Skip to content

fix(megatron): gather embedding evaluation metrics across DP ranks - #10091

Open
MrCapricornLiu wants to merge 2 commits into
modelscope:mainfrom
MrCapricornLiu:lch/fix-megatron-eval-dp-metrics
Open

fix(megatron): gather embedding evaluation metrics across DP ranks#10091
MrCapricornLiu wants to merge 2 commits into
modelscope:mainfrom
MrCapricornLiu:lch/fix-megatron-eval-dp-metrics

Conversation

@MrCapricornLiu

Copy link
Copy Markdown
Contributor

PR type

  • Bug Fix
  • New Feature
  • Document Updates
  • More Models or Datasets Support

PR information

Megatron embedding evaluation accumulates predictions and labels locally, so each data-parallel rank computes different metrics over its own shard. Pearson/Spearman correlations cannot be recovered by averaging these local correlations. InfoNCE evaluation also misses other DP ranks' in-batch candidates.

Give the embedding metric the trainer's data-parallel group and gather predictions together with their labels before computing the existing metrics. Gathering the accumulated lists supports unequal shard sizes and a rank with no local updates. The accumulated state stays local, so repeated computation does not duplicate samples. The default/Hugging Face entry remains local.

This complements #10088, which handles reranker query metrics. Embedding correlations and cross-query InfoNCE candidates require the complete predictions rather than a reduction of per-query scores. This change does not modify the reranker implementation.

Experiment results

  • The new CPU regression fails on the original implementation and passes after the fix. Both tests pass, including four Gloo workers in independent groups, unequal sample counts, an empty peer, repeated computation, reset, singleton groups, and the local/HF entry.
  • Four H800 GPUs with real Megatron-Core 0.16.0 process groups and NCCL: DP4, TP2×DP2 and CP2×DP2. Paired and InfoNCE metrics match a complete-data reference with and without an empty first DP rank, on every worker.
  • Changed-file pre-commit checks pass.

The GPU check exercises metric collection and process groups, not a full pretrained-model evaluation. Gathering embeddings adds communication and per-rank memory proportional to the evaluation data; full-model performance, pipeline parallel evaluation, and multi-node execution were not tested.

Signed-off-by: Chenghao Liu <chliu@stu.pku.edu.cn>
Signed-off-by: Chenghao Liu <chliu@stu.pku.edu.cn>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant