Repository navigation
Conversation
… run ends Media written under --shared-media-dir was content-addressed but never deleted, so screenshots from every run accumulated on the shared filesystem (a crashed CUA run left 3.7 GB, and a 1080p round adds about 0.6 GB). Each run now writes under <dir>/<run_uuid>, and the training drivers remove that directory as their last teardown step, after the rollout executor and trainers are gone. The directory is keyed by run rather than by rollout step because a sample generated in one step can train in the next.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f7f0e286e8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| def release_run_media(args) -> None: | ||
| if (media_dir := run_media_dir(args)) is not None: | ||
| shutil.rmtree(media_dir, ignore_errors=True) |
There was a problem hiding this comment.
When a multimodal run saves rollout checkpoints or replay dumps, teardown now deletes files those artifacts still reference. FullyAsyncRolloutFn.save() persists pending and completed Sample objects, and save_debug_rollout_data() serializes sample.to_dict(); neither copies the referenced media into the artifact. After a normal shutdown or an exception that runs the disposer, resuming buffered samples or replaying a dump therefore encounters missing media paths. Rebuilding the dataset does not repair the absolute references already stored in those samples. Preserve media reachable from saved artifacts, or copy it into artifact-owned storage and rewrite the references before deleting the run directory.
Useful? React with 👍 / 👎.
- Refuse train data that carries both media references and dense multimodal_train_inputs, from the default or a custom converter: the trainer rebuilds the tensors from the references, so shipping both silently pays the hand-off the references replace. - Cast each sample's rebuilt floating tensors to the training dtype as they leave the processor. SGLang's processor returns fp32 pixels, and get_rollout_data holds every sample's tensors for the whole step: at 1080p with two screenshots a sample that was 95.6 MiB, and round 0 ran out of GPU memory on every H200 rank. The vision tower runs in bf16, so the cast changes nothing it sees. - Mark the whole sequence as an expanded prefix when prefill scoring sends media references in token space. The tokens were already expanded during rollout, so without mm_token_expansion_start_len SGLang would expand the placeholders again.
The trainer builds one processor from the settings every rollout engine shares, and SglangConfig.common_value only notices engine groups that disagree when the first training step reads them, after a whole rollout round. Read those settings during argument parsing so a group override of the image processor, its config or the model fails at launch. A custom inference engine provider also implies external engines whose server args Miles never sees, so refuse it alongside --rollout-external-engine-addrs.
f7f0e28 to
d809e5b
Compare
Layer P15 of the CUA upstreaming stack (the Rockdu sync). It is built and reviewed in celve, and goes upstream with the whole stack once every layer is finished. Depends on radixark/miles radixark#4272 (Kangrui Du / Rockdu, branch
rockdu/vlm-02-token-spaceatf9e7a9f645in radixark/miles, stacked on radixark#4240rockdu/vlm-01-media-refsat3080d5fb1c), which we expect to land. This PR builds on his live head and carries none of his commits. Its celve base is a ref-for-ref mirror of his radixark branch, needed only because a celve PR can't use a radixark branch as base. The mirror is never edited, and it is re-synced, with this stack rebased, whenever radixark#4272 moves. Upstream, this PR opens against radixark'srockdu/vlm-02-token-spaceand is retargeted to main when radixark#4272 merges. It follows his 2026-10-29 SGLang token-space API. Issue: LIN-2519, hub LIN-2504.These are the CUA fixes we ran on top of a backport of his stack (#11), ported onto his current code. There are three commits, one per concern.
Problem.
--shared-media-diris content-addressed and never deleted. Screenshots from every run pile up on the shared filesystem: one crashed CUA run left 3.7 GB, and each 1080p round adds about 0.6 GB.multimodal_train_inputs, for example from a sample builder that still sets the tensors. When that happens it silently pays the hand-off that references replace, and the trainer then rebuilds the same tensors anyway.get_rollout_dataholds every sample's rebuilt tensors on the GPU for the whole step. At 1080p with two screenshots that is 95.6 MiB a sample. Round 0 of a CUA run ran out of GPU memory on every H200 rank, at 134.8 of 139.8 GiB.--recompute-logprobs-via-prefill) sends media references but nomm_token_expansion_start_len. In token space the tokens were already expanded during rollout, so SGLang would expand the placeholders again.RolloutEngineProcessorArgs.from_sglang_config).SglangConfig.common_valueonly finds engine groups that disagree when the first training step reads those settings, which is after a whole rollout round. Separately, a user-set--custom-inference-engine-provider-pathmeans external engines (_compute_rollout_external), and the startup check does not refuse it.Fix.
run_media_dir(args)returns<--shared-media-dir>/<run_uuid>. The dataset, eval and geo3k multi-turn call sites write there.release_run_media(args)removes that directory.train.pyandtrain_async.pyregister it first on the disposer, so it runs last, after the rollout executor and the trainers are gone. The directory is keyed by run, not by rollout step, because a sample generated in one step can train in the next. The--shared-media-dirhelp andcli-reference.mdsay so.convert_samples_to_train_dataraises when its output carries bothmultimodal_inputsandmultimodal_train_inputs. This applies to the default converter and to a custom one.compute_multimodal_train_inputstakes the training dtype and casts each sample's floating tensors as they leave the processor, in the per-sample rebuild both the round rebuild and a later per-micro-batch rebuild share.get_rollout_datapasses bf16, or fp16 under--fp16; fp32 runs are left as they are. The vision tower runs in the training dtype, so the cast changes nothing it sees._build_prefill_scoring_payloadsetsmm_token_expansion_start_len = len(sample.tokens)._check_multimodal_token_space_processing_argsnow callsRolloutEngineProcessorArgs.from_sglang_configwhile arguments are parsed, and turns a disagreement into aValueErrornaming the field. It also refuses external engines by_compute_rollout_external, not only by--rollout-external-engine-addrs. Thecli-reference.mdrow lists both.Dropped from the origin commits because Rockdu's
9de9319b83andbbddd9e654already cover them:validate_token_space_engine_settingsinlaunch_settings.py. That module does not exist on this branch.autoimage backend and the base--sglang-mm-process-config. The trainer now follows the engines' own settings.enable_token_space_processorcheck (c95dc47af9).common_value("enable_token_space_processor")covers it.--rollout-external-engine-addrsrefusal.Also dropped:
cf2eb9df40), which only measure.can_precompute_dp_schedulechange ina79b5f04fb, which belongs to our integration branch's dense-image schedule (Schedule image batches by rollout, Kimi and Inkling included #6) and not to this stack.Tests (CPU devbox; SGLang
sglang-milesb0ed761650, which has the 10-29 token-space API). This base predates radixark/miles radixark#4548, so itschat_template_utils/deepseek.pyimportsencoding_dsv4/encoding_dsv41from SGLang, andb0ed761650has removed them. To get past that, both modules were taken from SGLange22a530859^and appended to the package path through a test-onlysitecustomize. No code in this PR changes for it.f9e7a9f645(9 failed, 8 passed) and passes with the fix (17 passed). Failures without the fix:test_train.py::TestRunMediadrivestrain()and fails withassert not True: the run's directory is still there. It also checks that the media still exists when the rollout executor is disposed, and that another run's directory is kept.test_train_data_conversion.pyhas two tests, for the default and a custom converter. Both fail withDID NOT RAISE ValueError.test_data.py::TestGetRolloutDataMediaRebuilddrivesget_rollout_datathrough the realcompute_multimodal_train_inputs, with only the SGLang processor faked to return fp32 pixels. Onf9e7a9f645it fails withassert torch.float32 == torch.bfloat16and== torch.float16(2 failed, 1 passed). The fp32 case and the integerimage_grid_thwstay unchanged.test_prefill_logprobs.py::test_prefill_scoring_of_media_marks_the_expanded_tokens_as_prefix[True]fails withKeyError: 'mm_token_expansion_start_len'.test_arguments.py::TestMultimodalTokenSpaceProcessingArgs: the custom provider test and the two group-disagreement tests (anmm_process_configoverride and animage_processor_backendoverride) fail withDID NOT RAISE ValueError.pytest tests/fast/utils/test_media_refs.py tests/fast/test_train.py tests/fast/test_train_async.py tests/fast/backends/training_utils/test_data.py tests/fast/ray/rollout/test_train_data_conversion.py tests/fast/rollout/generate_utils/test_prefill_logprobs.py tests/fast/backends/training_utils/data/test_token_space_media_process.py tests/fast/utils/test_arguments.py: 657 passed. After the cast moved into the per-sample rebuild (headd809e5bf9c),pytest tests/fast/backends/training_utils/test_data.py tests/fast/backends/training_utils/data tests/fast/utils/test_arguments.py tests/fast/rollout/generate_utils/test_prefill_logprobs.py tests/fast/ray/rollout/test_train_data_conversion.py: 613 passed.pytest tests/fast/rollout/generate_utils tests/fast/rollout/test_data_source.py tests/fast/rollout/test_sglang_rollout.py tests/fast/rollout/inference_rollout/test_inference_rollout_common.py tests/fast/rollout/inference_rollout/test_compatibility.py tests/fast/utils/test_data.py tests/fast/ray/rollout/test_train_data_conversion.py(the changed call sites): 343 passed.pre-commit runon the changed files: passed.Limits.
train.pyandtrain_async.pyrelease run media.train_multi_policy.pyandserve_tinker.pydo not.Related upstream work.
miles/utils/media_refs.py, but in a different place (materialize_message_media), so the two do not conflict textually. ItsSessionConfigreadsargs.shared_media_dirdirectly. Once both land, it should readrun_media_dir(args), or its screenshots will sit outside the run's directory and never be removed.Origin commits (#11, on
LIN-1985/investigate-existing-recipes-for):1d5956e99d,a79b5f04fb,d2a5e6d631,6d10f16ef4,ad65487654,c95dc47af9(cf2eb9df40dropped).