Skip to content

fix: preserve per-sample training metadata across DP transport - #2403

Open
heyxingxing wants to merge 1 commit into
THUDM:mainfrom
heyxingxing:fix/preserve-rollout-train-metadata
Open

heyxingxing wants to merge 1 commit into
THUDM:mainfrom
heyxingxing:fix/preserve-rollout-train-metadata

Conversation

@heyxingxing

Copy link
Copy Markdown

Per-sample training metadata disappears before reaching the trainer: conversion ignores it when only later samples provide it, and DP transport omits the field even when conversion populated it. Include metadata when any sample provides it, retain None entries, and partition it alongside the other per-sample fields.

Closes #2400

Validation: PYTHONPATH=. python tests/test_process_rollout_data.py20 passed. Twelve added regression cases cover DP sizes 1/2, static/dynamic batching, and complete/mixed/absent metadata; eight of them fail on the base implementation with missing metadata. Tests exercise the real conversion, scheduler and consumer with serving/telemetry imports and Ray object-store calls isolated. The existing CPU CI registration is reused.

Ruff 0.14.7, Black 24.3.0, isort 5.13.2 and git diff --check passed. CPU validation only (Python 3.12.13, macOS arm64); no GPU or multi-process Ray run.

This branch has not been deployed

No deployments
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.

[Bug] Per-sample train_metadata is lost during conversion and DP transport

1 participant