Skip to content

Offload debug rollout dump to a Ray task; drop routed_experts, add --save-rollout-interval - #83

Draft
nightlessbaron wants to merge 1 commit into
prodfrom
feat/async-debug-rollout-dump
Draft

Offload debug rollout dump to a Ray task; drop routed_experts, add --save-rollout-interval#83
nightlessbaron wants to merge 1 commit into
prodfrom
feat/async-debug-rollout-dump

Conversation

@nightlessbaron

Copy link
Copy Markdown

--save-debug-rollout-data serialized the full sample list with torch.save on the RolloutManager MainThread. With routed_experts in every sample the dump is ~200GB: ~30 minutes holding the GIL (frozen event loop, idle engines, blocked training handoff) and a peak-memory doubling that can OOM the head node.

Three changes:

  • The write runs in a @ray.remote(num_cpus=1) task — its own worker process, so serialization never blocks the actor. At most one write is in flight; an overlapping dump is skipped with a warning instead of queued.
  • rollout_routed_experts is dropped from the dump (it dominates the size and the training path already carries it), shrinking the payload ~40x and keeping the object-store handoff cheap.
  • New --save-rollout-interval N dumps only every Nth training rollout (default 1; eval dumps unaffected).

Known trade-offs: the final rollout's dump is fire-and-forget, so a job that exits immediately after its last step can leave that one file incomplete; retention's walk-backward delete pairs with interval=1 (gaps from larger intervals can stop it early).

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