Conversation
Provides DumperPhase type, CLI config parsing with key validation, phase-specific directory helpers, env var builder for SGLang subprocess, and dumper_phase_scope context manager for Megatron lifecycle.
Adds --dumper-enable, --dumper-dir, --dumper-sglang, --dumper-fwd-only, --dumper-fwd-bwd flags. Links --dump-details to auto-enable dumper.
Integrates sglang dumper lifecycle into Megatron forward-only and forward-backward passes via dumper_phase_scope context manager.
Each engine gets isolated dump directory via DUMPER_EXP_NAME=engine_{rank}.
Runs a full training loop with --dumper-enable, then verifies dump files exist in all three phase directories including per-engine isolation for SGLang.
Unified naming eliminates _PHASE_ATTR_MAP dict — CLI attr is now
derived as f"dumper_{phase}". Renamed --dumper-sglang to
--dumper-inference for consistency.
Eliminates separate PHASE_* constants and _ALL_PHASES tuple. Callers use DumperPhase.FWD_ONLY etc. directly.
- parse_dumper_config: use type-aware coercion via _FrozenConfig._parse_env_value instead of value-based heuristics (e.g. "0" no longer becomes False for str fields) - configure_dumper_for_phase: reset to defaults before each phase to prevent cross-phase config leaking via singleton replace() - dumper_phase_scope: disable dumper on scope exit to prevent stray dumps - sglang_engine: clean up DUMPER_* env vars after server launch - Fix --dumper-enable help text referencing old --dumper-sglang flag
Remove parse_dumper_config from dumper_utils.py — the parsing logic now lives in _DumperConfig (sglang side) so other users can reuse it. Internally use _kv_pairs_to_dict for partial-override semantics.
Add extra_env parameter to launch_server_process that sets env vars inside the child process only, avoiding the os.environ.update/pop pattern that leaked state in the parent Ray actor.
- Remove dumper_phase_scope, add finalize_dumper_phase - Call configure_dumper_for_phase before and finalize_dumper_phase after - Move step() to end (after dump_model) instead of beginning
Instead of passing env vars through extra_env/multiprocessing wrapper in sglang_engine.py, inject DUMPER_* vars into the Ray actor's runtime_env alongside other env vars. The spawned SGLang server subprocess inherits them naturally.
Instead of pre-setting all DUMPER_* env vars at Ray actor creation,
only set DUMPER_SERVER_PORT=reuse to register the HTTP endpoint.
Configure and enable the dumper dynamically via POST /dumper/configure
before each rollout, giving each engine its own exp_name=engine_{i}.
Also extract _get_worker_urls helper in sglang_rollout.py.
Callers no longer need to obtain and pass worker_urls — the function discovers them itself via get_worker_urls from inference_rollout_train.
Inline worker URL discovery back into abort() to keep sglang_rollout.py diff minimal. Enable cleanup_previous by default for both SGLang and Megatron dumper phases.
dir = dumper_dir for all phases, exp_name carries the distinction:
engine_{i} for inference, fwd_only / fwd_bwd for Megatron.
Eliminates redundant fwd_only/fwd_only nesting.
Wrap forward_step with dumper stepping so each microbatch gets its own dumper step counter. The wrapper calls dumper.step() before every forward_step invocation except the first, ensuring model forward and the subsequent loss callback share the same step number.
Module name already conveys the dumper context. Callers now use `dumper_utils.configure_for_phase(...)` instead of importing `configure_dumper_for_phase`. Also switch all call sites from `from dumper_utils import X` to `from miles.utils import dumper_utils`.
Prevents health-check requests from polluting dumper output by overriding three args in miles_validate_args: - use_fault_tolerance=False (RolloutHealthMonitor) - router_disable_health_check=True (sgl-router) - rollout_health_check_interval=1e18 (miles-router)
Callers were passing model[0] which would silently lose parameters when virtual pipeline parallelism is enabled (len(model) > 1). Now finalize() receives the full model list and asserts len == 1 until multi-chunk support is implemented.
…tensors dump_model() saved every parameter as a separate .pt file, which for MoE models (e.g. 128 experts) produced files 40+ GB each, totaling 500+ GB across ranks. Replace with save=False (console only) plus a single lightweight model_summary dump containing shapes/dtypes/devices.
…of full tensors" This reverts commit 5e2f648.
Speeds up e2e test by generating shorter responses (128 -> 20 tokens).
With enable_model_value and enable_model_grad defaulting to False, the test must explicitly enable them via phase overrides so that fwd_only and fwd_bwd directories are created for verification.
MoE model has too many parameters for dump_model to complete in reasonable time. Explicitly disable enable_model_value/grad and only verify engine_* (inference) dumps.
The non-intrusive dumper (non_intrusive_mode=core) needs register_non_intrusive_dumper() to be called with the model so forward hooks can dump input_ids, positions, etc. Without this, fwd_only and fwd_bwd phases produce no dump files. Restore EXP_PATTERNS to verify all three phases.
Pass model to DumperMegatronUtil.__init__ and register non-intrusive hooks there. This is safe because _configure() calls dumper.reset() which now removes previous hooks before re-registering. Removes the separate register_hooks() method and its call sites in forward_only() and train_one_step().
Megatron models pass input_ids/positions as keyword arguments,
which non-intrusive hooks cannot capture via pre_hook positional
args. Core mode only dumps fields matching {input_ids, positions},
so it produces no output for Megatron. Switch to 'all' mode to
capture sub-module I/O via positional args.
The barrier in sglang's cleanup_previous deadlocks under pipeline parallelism. Miles handles dump directory cleanup externally via prepare() before training, so the lazy barrier-based cleanup is not needed.
non_intrusive_mode=all on a 30B MoE model generates thousands of dump files per forward pass, causing excessive I/O. Filter to only dump embedding-related modules to verify hooks work without overwhelming disk.
- Remove cleanup_previous from sglang dumper config to avoid dist.barrier() deadlocks in async PP contexts - Add _cleanup_dump_dir helper: rank-0 rmtree + barrier - Call cleanup explicitly in DumperMegatronUtil._configure - E2e test: switch from non_intrusive_mode=all to default core mode, disable model value/grad dumps for both fwd phases
The DumperConfig default is already False, and no caller passes this field explicitly.
Avoids rmtree on files or symlinks to files.
Move torch.distributed import to module level, extract _get_rank helper, flatten the if/elif into two independent checks.
Check that specific dump fields exist for each phase: - engine_*: input_ids, positions (from SGLang ForwardBatch) - fwd_only/fwd_bwd: input_ids, cu_seqlens_q, cu_seqlens_kv, qkv_format (from Megatron kwargs + PackedSeqParams)
Dump files are named like step=0___rank=0___name=input_ids.pt,
not just input_ids.pt. Use *name={field}.pt glob pattern.
This reverts commit d3a82dc.
fzyzcjy
requested review from
guapisolo,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
Summary of ChangesHello @fzyzcjy, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a robust dumping mechanism designed to aid in the detailed analysis and debugging of both SGLang inference and Megatron training workflows. It provides granular control over data collection through new command-line arguments and intelligently adjusts runtime configurations to prevent interference with the dumping process, thereby enhancing the observability of complex model behaviors. Highlights
Changelog
Ignored Files
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a "dumper" utility to capture intermediate states during different phases of training and inference for debugging and analysis. The feature is enabled by a new set of --dumper-* command-line arguments. When enabled, it overrides several other arguments to ensure a controlled environment for dumping, such as disabling health checks and forcing a single rollout. The implementation is well-structured and includes new utility functions, integration into Megatron and SGLang workflows, and comprehensive unit and end-to-end tests.
My review focuses on improving code structure and readability. I've pointed out a circular dependency that could be refactored and suggested replacing a magic number with a more idiomatic constant.
| if not _is_phase_enabled(args, DumperPhase.INFERENCE): | ||
| return | ||
|
|
||
| from miles.rollout.inference_rollout.inference_rollout_train import get_worker_urls |
There was a problem hiding this comment.
There appears to be a circular dependency here. miles/utils/dumper_utils.py imports from miles/rollout/inference_rollout/inference_rollout_train.py, which in turn imports miles/utils/dumper_utils.py. While using a local import inside configure_sglang makes this work at runtime, it's generally better to avoid circular dependencies for improved code structure and maintainability.
Consider moving the get_worker_urls function from miles/rollout/inference_rollout/inference_rollout_train.py to a lower-level utility module like miles/utils/http_utils.py, since it's already using get from there. Both dumper_utils.py and inference_rollout_train.py could then import it from http_utils.py without creating a cycle.
|
|
||
| logger.info("Dumper mode: all heartbeat mechanisms disabled") | ||
| args.router_disable_health_check = True | ||
| args.rollout_health_check_interval = 1e18 |
There was a problem hiding this comment.
Using 1e18 to effectively disable the health check interval works, but it's a bit of a magic number. For better readability and to more clearly express the intent of an infinite timeout, consider using float('inf').
| args.rollout_health_check_interval = 1e18 | |
| args.rollout_health_check_interval = float('inf') |
|
|
||
| assert args.use_fault_tolerance is False | ||
| assert args.router_disable_health_check is True | ||
| assert args.rollout_health_check_interval == 1e18 |
There was a problem hiding this comment.
To align with the suggested change of using float('inf') for the health check interval in miles/utils/arguments.py, this assertion should be updated accordingly.
| assert args.rollout_health_check_interval == 1e18 | |
| assert args.rollout_health_check_interval == float('inf') |
yushengsu-thu
left a comment
yushengsu-thu
left a comment
There was a problem hiding this comment.
workflow part LGTM
2 participants
No description provided.