Make the tensor dumper fault-tolerance aware - #1424
Merged
fzyzcjy merged 42 commits intoJul 10, 2026
Merged
Conversation
fzyzcjy
requested review from
Shi-Dong,
Zhichenzzz,
guapisolo,
jybsuper,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
June 22, 2026 10:38
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces rollout_id propagation to organize dump directories per rollout and adds support for gathering distributed-optimizer grad shards. It also ensures dump files are only written on DP rank 0 and handles NFS cleanup errors gracefully. The review feedback highlights a critical issue where wiping the entire phase directory at rollout_id == 0 can delete active dumps from other surviving cells in multi-cell or fault-tolerant environments, and provides a code suggestion to gate this cleanup.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+164
to
+165
| if rollout_id == 0: | ||
| _cleanup_dump_dir(Path(merged["dir"]) / phase.value) |
Contributor
There was a problem hiding this comment.
In a multi-cell (independent-DP) or fault-tolerant training environment, multiple cells can start or restart asynchronously. If a cell starts or joins later and triggers rollout_id == 0, wiping the entire phase.value directory will delete the active dumps already written by other surviving cells. We should gate this full directory cleanup to only run when independent DP or the experimental FT trainer is not enabled.
Suggested change
| if rollout_id == 0: | |
| _cleanup_dump_dir(Path(merged["dir"]) / phase.value) | |
| if rollout_id == 0 and not (getattr(args, "indep_dp", False) or enable_experimental_ft_trainer()): | |
| _cleanup_dump_dir(Path(merged["dir"]) / phase.value) |
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/thread-witness-ids-through-the-training-data-path
branch
from
June 23, 2026 07:49
bfc1ec7 to
bd0ed91
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/make-the-tensor-dumper-fault-tolerance-aware
branch
from
June 23, 2026 07:49
737dfbf to
d82ff55
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/thread-witness-ids-through-the-training-data-path
branch
from
June 23, 2026 09:27
bd0ed91 to
cd1318d
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/make-the-tensor-dumper-fault-tolerance-aware
branch
from
June 23, 2026 09:28
d82ff55 to
b1d4e89
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/thread-witness-ids-through-the-training-data-path
branch
from
June 23, 2026 13:31
cd1318d to
720c191
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/make-the-tensor-dumper-fault-tolerance-aware
branch
from
June 23, 2026 13:31
b1d4e89 to
4dd9947
Compare
| if not overrides.get("enable"): | ||
| return False | ||
|
|
||
| exp_name = f"{phase.value}/rollout_{rollout_id}" |
Collaborator
There was a problem hiding this comment.
maybe add store_prefix (actor/critic/teacher...) here to avoid conflict
yueming-yuan
approved these changes
Jun 29, 2026
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/thread-witness-ids-through-the-training-data-path
branch
from
July 8, 2026 03:54
720c191 to
e2c968f
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/make-the-tensor-dumper-fault-tolerance-aware
branch
from
July 8, 2026 03:54
4dd9947 to
b2b8526
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/thread-witness-ids-through-the-training-data-path
branch
from
July 8, 2026 05:56
e2c968f to
68062b7
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/make-the-tensor-dumper-fault-tolerance-aware
branch
from
July 8, 2026 05:56
b2b8526 to
078092f
Compare
Add a `deterministic_random` reward that hashes the sample tokens + response to produce a stable pseudo-random 0/1 reward, used for reproducible fault-tolerance / CI tests. - rm_hub/__init__.py (+ test).
Add `inplace_modify_args`, a context manager that temporarily overrides args attributes and restores them on exit (asserting they weren't clobbered), used to scope per-attempt argument overrides in the fault-tolerant trainer. - argparse_utils.py (+ test).
Small shared-utility additions used by the fault-tolerant trainer: hash non-contiguous tensors safely (reshape before viewing as bytes), an `enable_experimental_ft_trainer` env flag, forward NCCL_DEBUG/NCCL_DEBUG_FILE to worker environments, and a `filter_keys` helper. - ci_utils.py / environ.py / external_utils/command_utils.py / misc.py.
Thread the original backend through ReloadableProcessGroup so that, when a process group is rebuilt (e.g. after a reconfigure/heal), it is recreated with the same backend instead of hard-coding NCCL. - reloadable_process_group.py: carry `backend` in the reload group info.
Add small foundation utilities used across the fault-tolerance trainer: a strict pydantic base model, a retry helper, a tensor checksum helper, a per-cell megatron world-size computation, the TrainStepOutcome enum, and the IndepDPInfo dataclass describing a cell's independent-DP identity. - pydantic_utils.py / retry_utils.py / checksum_utils.py / megatron_args_utils.py / types.py / indep_dp.py and tests.
Add a `log_structured` helper that emits logfmt-style key/value log lines, used by the fault-tolerance components for greppable structured logs. - structured_log.py (+ test).
Add a small `Clock` interface (`RealClock` plus a controllable fake clock) so time-dependent fault-tolerance code (health checks, heartbeats) can be driven deterministically in tests. - miles/utils/clock.py and tests.
Add a fault-injector test utility used to deterministically exercise fault-tolerance code paths. - miles/utils/test_utils/fault_injector.py.
Add the shared data models for the fault-tolerance control server (e.g. the `TriState` health value), used by the health checker and later by the HTTP control server. - miles/utils/control_server/models.py.
Add the structured event logger that records typed events keyed by per-process identity, wire it through the logging helper and CLI argument, and start it from the train entrypoints. - miles/utils/event_logger/logger.py, logging_utils.py, arguments.py and entrypoint wiring, with tests.
Add snapshot/restore for the structured event log so the event history survives cell restarts during fault-tolerant training. - miles/utils/event_logger/checkpoint.py and tests.
Add the `MetricEvent` model (a discriminated-union member) and emit every
tracking metric into the structured event log: `tracking_utils.log` now forwards
`{metrics}` to `get_event_logger().log(MetricEvent, ...)` when the event logger
is initialized.
Add the witness id allocator and `WitnessInfo` carrier used to assign and track witness ids for fault-tolerance verification. - miles/utils/witness/allocator.py and tests.
Thread witness ids through the model by injecting witness parameters, so the event log can later verify they propagate correctly. - miles/utils/witness/module.py, model_provider.py and tests.
Add the first event-analyzer rules that replay the structured event log and flag
weight-checksum inconsistencies: a `checksum_compare` helper (flatten nested
dicts, diff flat checksum maps) plus two rules built on it — cross-replica weight
checksum consistency and inference-engine weight checksum consistency — with unit
tests.
- miles/utils/event_analyzer/rules/{checksum_compare,cross_replica_weight_checksum,inference_engine_weight_checksum_consistency}.py and tests.
Add the witness-tracing rule for the event analyzer: it follows witness ids through the replayed event log to verify they are propagated correctly across the training pipeline, with unit tests. - miles/utils/event_analyzer/rules/witness.py and tests.
Add the analyzer that replays the structured event log and applies the analysis rules (checksum-consistency and witness tracing) to verify fault-tolerance behaviour offline, with unit tests. - miles/utils/event_analyzer/analyzer.py and tests.
Add comparison helpers used by fault-tolerance tests to compare dumped tensors and
inference-engine checksums offline (generic comparators, dump comparison, and an
inference-engine checksum comparison built on the event-analyzer checksum rule).
- miles/utils/test_utils/comparisons/{comparators,dumps,inference_engine_checksums}.py and tests.
Add the metric comparison helpers used by fault-tolerance tests to compare logged training metrics offline. - miles/utils/test_utils/comparisons/metrics.py.
Add reusable reconfiguration assertions used by fault-tolerance tests to verify cell reconfigure / healing behaviour, with unit tests. - miles/utils/test_utils/reconfigure_assertions.py and tests.
Move the GroupInfo dataclass out of training_utils/parallel.py into a shared process_group_utils module that also provides multi-process-group helpers (GeneralPGUtil / MultiPGUtil) for collectives over single or hierarchical process groups, used by the cross-replica / effective-DP code paths. - process_group_utils.py: GroupInfo + GeneralPGUtil / MultiPGUtil (+ tests). - training_utils/parallel.py: import GroupInfo from the shared module. - megatron / fsdp parallel.py: import GroupInfo from the shared module. - distributed_utils.py: use GeneralPGUtil for masked-whiten all-reduce.
Introduce `_TensorViewCodec` (in `checkpoint_transfer.py`), which encodes a list of tensors as (unique_storages, view_metas) by deduping shared underlying storages (e.g. Megatron distributed-optimizer grad buckets) and reconstructs the original views via `as_strided`. Comprehensively unit-tested by `TestTensorViewCodec`. Consumed by the peer checkpoint transfer added next.
Add an InMemoryCheckpointManager that keeps checkpoints in memory and thread a `checkpointing_context` / `non_persistent_ckpt` through the model save path so a non-persistent checkpoint can be requested. - miles/backends/megatron_utils/in_memory_checkpoint.py, model.py and tests.
Add peer-to-peer checkpoint transfer (built on the in-memory checkpoint manager) so a healed cell can receive weights from a surviving peer. - miles/backends/megatron_utils/checkpoint_transfer.py and tests.
Move the inline env-var / backend-selection / ray.remote actor construction out of RayTrainGroup into a module-level allocate_gpus_for_actor factory, with no behaviour change. (Pure mechanical extraction; FT context is threaded in a follow-up.) - miles/ray/train/actor_factory.py, miles/ray/actor_group.py.
…hrough the train group Add the FT context to the actor factory and train actors: indep_dp_store_addr / role / cell_index params and per-process-identity logging in TrainRayActor; pass a rollout_manager into RayTrainGroup (set_rollout_manager, pause rollout health monitors before weight updates); freeze the v1 group and wire context through placement_group / entrypoints.
Add the RayTrainCell abstraction (a group of actors forming one independent-DP replica) together with its cell state model, with the test harness. - miles/ray/train/cell.py, miles/ray/train/cell_state.py and tests.
…al flag Add the cell-based independent-DP train group (RayTrainGroupV2) that orchestrates RayTrainCells, selected via the experimental flag, and dispatch to it from the placement group. - miles/ray/train/group.py, placement_group.py, actor_group.py and actor wiring, with tests.
…nRayActor Also tidy formatting in the Megatron actor and model touched here.
When independent DP is enabled the megatron data-parallel size is 1 per cell, so omit the `dp_size` hint from the actor's train parallel config (cross-replica DP is handled separately) instead of reporting the intra-cell dp size. - actor.py: gate `train_parallel_config` on `args.indep_dp`.
Pull the bshd/thd slice-pad-stack transform out of the position-ids path in get_batch into a local _compute_transform_like_token_ids helper, so the same transform can be reused for other per-token id streams. No behaviour change. - miles/backends/training_utils/data.py.
Carry per-sample witness ids from the rollout data into training so each token can be tagged with its originating witness id (used by the event analyzer to trace samples across rollouts), reusing the CP-aware token-id transform helper in get_batch. - utils/data.py / training_utils/data.py: attach seq_witness_ids and expand into per-token witness_ids. - actor.py / model.py: thread witness_info / attempt through train / train_actor / train_one_step. - log_utils.py: exclude witness_ids from rollout-data logging.
Scope tensor dumps per rollout and make them safe under independent-DP cells: dumps go into per-rollout subdirectories, only effective-DP rank 0 writes output files (other ranks still join dumper collectives), weights/grads are dumped once per rollout pinned to step 0 with a distributed-optimizer grad all-gather, and dump-dir cleanup is resilient to crashed peers and barriers across cells. - dumper_utils.py: per-rollout exp_name, rank-gated output, `_build_full_grad_getter`, resilient `_cleanup_dump_dir` + `_barrier_after_dump_dir_cleanup`. - model.py / actor.py: thread `rollout_id` into `forward_only` / `compute_log_prob` and the DumperMegatronUtil construction.
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/thread-witness-ids-through-the-training-data-path
branch
from
July 10, 2026 02:12
68062b7 to
af73984
Compare
fzyzcjy
force-pushed
the
tom/pr_chain/trainer_ft/dev_revert_reversed/make-the-tensor-dumper-fault-tolerance-aware
branch
from
July 10, 2026 02:12
078092f to
c5f3c25
Compare
Base automatically changed from
tom/pr_chain/trainer_ft/dev_revert_reversed/thread-witness-ids-through-the-training-data-path
to
main
July 10, 2026 03:23
…ft/dev_revert_reversed/make-the-tensor-dumper-fault-tolerance-aware
fzyzcjy
deleted the
tom/pr_chain/trainer_ft/dev_revert_reversed/make-the-tensor-dumper-fault-tolerance-aware
branch
July 10, 2026 03:24
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Scope tensor dumps per rollout and make them safe under independent-DP cells:
dumps go into per-rollout subdirectories, only effective-DP rank 0 writes output
files (other ranks still join dumper collectives), weights/grads are dumped once
per rollout pinned to step 0 with a distributed-optimizer grad all-gather, and
dump-dir cleanup is resilient to crashed peers and barriers across cells.
_build_full_grad_getter,resilient
_cleanup_dump_dir+_barrier_after_dump_dir_cleanup.rollout_idintoforward_only/compute_log_proband the DumperMegatronUtil construction.