feat: add bounded asynchronous hidden-state fan-out#811
Conversation
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe PR adds bounded shared hidden-state fan-out for independent trainers, a benchmark harness with GPU telemetry and request accounting, and optional DFlash/Liger fused training paths. It also adds CLI wiring, optimizer controls, extensive validation, and unit/integration coverage. ChangesShared hidden-state fan-out
Independent-consumer benchmark
DFlash and optimizer paths
Possibly related PRs
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (2)
tests/unit/models/test_dflash_optimized_ce.py (1)
72-103: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd draft-vocabulary coverage for both label sources.
Verify that
input_idsis rejected whenuse_draft_vocab=True, while verifier-argmax targets remain valid reduced-vocabulary IDs without another remapping step.As per path instructions, tests should cover vocabulary-mapping boundaries. Based on learnings, sliced verifier-head logits are already in draft-vocabulary space.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/models/test_dflash_optimized_ce.py` around lines 72 - 103, Add tests around _ce_target_ids covering both label_source values with use_draft_vocab=True: assert label_source="input_ids" is rejected, and assert verifier-argmax targets remain valid draft-vocabulary IDs directly from the sliced verifier-head logits without additional remapping. Include boundary cases for the vocabulary mapping.Sources: Path instructions, Learnings
tests/unit/benchmarks/test_independent_consumers.py (1)
135-152: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winValidate the example’s windowed multiplicity contract.
This test only inspects
--train-data-ratio, so the contradictory 1P3C multiplicity remains undetected. Load the fixture throughload_configand assert that a windowed multi-consumer scenario uses expected service multiplicity1.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/benchmarks/test_independent_consumers.py` around lines 135 - 152, Update test_example_config_uses_train_only_workloads to load the fixture via load_config instead of only json.loads, then inspect the windowed multi-consumer scenario and assert its expected service multiplicity is 1. Preserve the existing command count and --train-data-ratio assertions while adding coverage for the 1P3C multiplicity contract.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@scripts/benchmark_independent_consumers.py`:
- Around line 18-26: Update the argument parser for --run-directory and --report
so they are optional during --validate-only execution, while still requiring
both when launching a benchmark. Adjust the validation or main execution flow to
enforce these paths only when validate-only is not set, preserving their
existing Path types and behavior.
In `@src/speculators/benchmarks/gpu_monitor.py`:
- Around line 431-463: The GPU window aggregation in the per-GPU summary
currently omits compute-process memory. Extend the `gpu_rows` aggregation to
collect non-null `used_memory_bytes` values from each row’s `compute_processes`,
then expose their maximum as a separate peak role-memory field in `per_gpu`,
using the existing MiB conversion and `None` behavior when no values are
present. Keep total device-memory aggregation unchanged.
- Around line 437-450: Update the per-GPU validation logic around gpu_rows and
_MIN_COVERAGE_SAMPLES so windows with fewer than two samples are marked invalid
and recorded in invalid_reasons. Ensure the existing valid path remains
unchanged for windows meeting the minimum sample requirement, and do not treat a
single sample with zero coverage as reliable evidence.
- Around line 228-256: Update GpuMonitor.start so the entire initialization
sequence, including the session-start _write and _thread.start calls, is covered
by failure cleanup rather than only backend.open. On any exception, close and
clear _output and invoke _close_backend(), while preserving the original
exception and successful startup behavior.
In `@src/speculators/benchmarks/independent_consumers.py`:
- Around line 91-103: Update _validate_single_process_command to reject
distributed launcher module names invoked via python -m, including
torch.distributed.launch and accelerate.commands.launch, alongside the existing
torch.distributed.run check. Preserve acceptance of direct Python commands that
do not invoke a distributed launcher.
- Around line 1169-1172: Update the windowed_artifacts_enabled detection over
scenario.consumers so each command token is normalized by splitting once at “=”
before comparing with --shared-hidden-states-consumer-id. Preserve recognition
of the standalone option while also enabling windowed mode for
--shared-hidden-states-consumer-id=value.
- Around line 1087-1099: Update terminate to catch and suppress
ProcessLookupError around both os.killpg calls, including the SIGTERM and
timeout-triggered SIGKILL paths. Continue waiting on self.process after either
race so the child is reaped, then preserve the existing finished_at and
close_log cleanup.
- Around line 193-200: Enforce the bounded-window multiplicity contract across
all sites: in src/speculators/benchmarks/independent_consumers.py lines 193-200,
update the model validation around
expected_service_completions_per_shared_sample to reject windowed multi-consumer
configurations unless multiplicity is 1; in
benchmarks/independent_consumer_fanout/config.example.json lines 312-318, change
the 1P3C expected multiplicity from 3 to 1; and in
tests/unit/benchmarks/test_independent_consumers.py lines 135-152, load the
fixture through the Pydantic model and assert this cross-field validation rule.
In `@src/speculators/data_generation/windowed_artifacts.py`:
- Around line 397-480: The register_positions flow must prune committed position
metadata so the positions table remains bounded across epochs. Update
register_positions and the corresponding completion path to delete positions
older than every active consumer’s lookbehind boundary, preserving positions
still needed by any consumer; add a multi-epoch test covering retention and
pruning behavior.
In `@src/speculators/train/data.py`:
- Around line 625-639: Update the pending-futures loops around wait in both
capture paths to renew each active claim while work remains pending and check
_windowed_producer_stop during each wait cycle. When shutdown is requested,
cancel pending futures, stop waiting for new results, and ensure the executor is
shut down without continuing background capture work; preserve
coordinator.fail_generation handling for completed futures and release or
finalize claims appropriately.
In `@src/speculators/train/dataloader.py`:
- Around line 272-276: Update the train and validation shared-artifacts consumer
ID construction near shared_artifacts_consumer_id so each data-parallel rank
receives a distinct coordinator key, incorporating the current dp_rank into the
existing :train and :val IDs. Apply the same change to both referenced branches
while preserving None when no consumer ID is configured.
In `@tests/unit/ops/test_fused_linear_cross_entropy.py`:
- Line 89: Update the pytest.raises match patterns in the affected tests to
escape the dot in the literal “torch.long” text, including both occurrences, so
Ruff RUF043 no longer treats the dots as regex metacharacters.
---
Nitpick comments:
In `@tests/unit/benchmarks/test_independent_consumers.py`:
- Around line 135-152: Update test_example_config_uses_train_only_workloads to
load the fixture via load_config instead of only json.loads, then inspect the
windowed multi-consumer scenario and assert its expected service multiplicity is
1. Preserve the existing command count and --train-data-ratio assertions while
adding coverage for the 1P3C multiplicity contract.
In `@tests/unit/models/test_dflash_optimized_ce.py`:
- Around line 72-103: Add tests around _ce_target_ids covering both label_source
values with use_draft_vocab=True: assert label_source="input_ids" is rejected,
and assert verifier-argmax targets remain valid draft-vocabulary IDs directly
from the sliced verifier-head logits without additional remapping. Include
boundary cases for the vocabulary mapping.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 33b98afb-923c-445d-bc7f-148d7c962e6a
📒 Files selected for processing (36)
benchmarks/independent_consumer_fanout/README.mdbenchmarks/independent_consumer_fanout/config.example.jsondocs/cli/train.mdpyproject.tomlscripts/benchmark_independent_consumers.pyscripts/train.pysrc/speculators/benchmarks/__init__.pysrc/speculators/benchmarks/gpu_monitor.pysrc/speculators/benchmarks/independent_consumers.pysrc/speculators/data_generation/artifact_cache.pysrc/speculators/data_generation/windowed_artifacts.pysrc/speculators/models/attention.pysrc/speculators/models/dflash/core.pysrc/speculators/models/dflash/metrics.pysrc/speculators/ops/__init__.pysrc/speculators/ops/fused_linear_cross_entropy.pysrc/speculators/train/data.pysrc/speculators/train/dataloader.pysrc/speculators/train/optimizers.pysrc/speculators/train/trainer.pytests/integration/models/test_model_forward.pytests/unit/benchmarks/__init__.pytests/unit/benchmarks/test_gpu_monitor.pytests/unit/benchmarks/test_independent_consumers.pytests/unit/data_generation/test_artifact_cache.pytests/unit/data_generation/test_windowed_artifacts.pytests/unit/models/test_attention.pytests/unit/models/test_dflash_metrics.pytests/unit/models/test_dflash_optimized_ce.pytests/unit/ops/__init__.pytests/unit/ops/test_fused_linear_cross_entropy.pytests/unit/ops/test_fused_linear_cross_entropy_cuda.pytests/unit/train/test_cli_args.pytests/unit/train/test_optimizers.pytests/unit/train/test_shared_artifacts.pytests/unit/train/test_windowed_training.py
|
Scope clarification: Mooncake is intentionally not included in this change. A related Mooncake-backed implementation is publicly available in The updated description now also links the existing upstream Mooncake work in #605/#710 |
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
|
Review findings are addressed in commit 967ddf5. The follow-up fixes make validation-only output paths optional; close partial NVML startup cleanly; distinguish device-wide from compute-process memory; reject one-sample, zero-coverage, and incomplete process-memory windows; reject additional distributed launchers; enforce publish-once multiplicity for windowed 1P3C; tolerate process-exit races; recognize option=value consumer IDs; bound position metadata across epochs; renew and release capture claims during shutdown; scope consumer IDs by DP rank; and cover the two review nits. Validation:
The report was valid with no invalid reasons; all role processes exited with code 0. This was a focused review validation, not a full-suite run. |
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
|
Follow-up validation after the issue #1 audit fixes Pushed commits:
The follow-up validation used the Python 3.12 project environment. No full test suite or CI was run for this follow-up. |
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
|
Reviewability and scope note: this PR is intentionally being left as the current end-to-end change for now, but I recognize that 37 files and roughly 11.5k additions make it substantially larger than a typical Speculators PR. Rather than rewrite the branch speculatively and discard existing review context, I will keep the branch maintained while waiting for maintainer direction. If maintainers are interested in the functionality but need smaller review units, I can split it into a dependency-ordered series covering: (1) the shared artifact store and coordinator, (2) bounded-window training and CLI integration, (3) the DFlash consumer performance path, and (4) benchmark/NVML validation. The split would preserve the tested end-to-end behavior and performance evidence. Mooncake remains deliberately outside #811. A follow-up can build on the transfer abstraction from #735 and coordinate with the active Mooncake work in #605/#710, instead of adding another backend and validation matrix to this already large review. |
|
Final status update at
This PR is ready for maintainer review. |
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
Design-level assessment: The implementation is well-structured -- opt-in behavior is properly gated, existing training paths are unchanged, validation is thorough, and the test suite is comprehensive (198+ focused tests). The POSIX filesystem constraints are clearly documented and the code does not overreach on NFS or GPU-direct claims. The THIRD_PARTY_NOTICES attribution is appropriate.
However, this PR has a scope problem that makes meaningful code review impractical.
Splittability: At 37 files and ~11,500 added lines, this bundles at least four independent feature areas:
- Shared artifact store and coordinator (
artifact_cache.py,windowed_artifacts.py) -- the core filesystem cache and SQLite-backed window coordination. - Bounded-window training integration (
data.py,dataloader.py,trainer.py,train.pyCLI wiring) -- lifecycle management, lease semantics, producer thread. - DFlash consumer performance path (
fused_linear_cross_entropy.py,dflash/core.py,metrics.py,optimizers.py) -- Liger CE backend, AdamW backend selection, fused gradient clipping. - Benchmark harness and NVML telemetry (
independent_consumers.py,gpu_monitor.py, benchmark fixtures) -- 1P1C/1P3C orchestration and evidence collection.
The DFlash/optimizer changes (#3) are fully independent of the windowed artifact system. The benchmark harness (#4) tests the windowed system but is a standalone deliverable. Even #1 and #2 could be separated into a store-only PR followed by a training-integration PR. The author has already offered this split -- I would recommend taking them up on it. Smaller PRs get deeper reviews and faster iteration.
CONTRIBUTING.md process: This PR resolves #810, but that issue is not assigned to any contributor. Per the project's contribution guidelines, significant changes (new features, CLI/API modifications, 3+ files) should get a maintainer "green light" via issue assignment before implementation begins. There are also no maintainer comments on #810. This is a process observation for maintainers, not a code concern.
Code quality: After detailed review of all source files, the code does not exhibit AI-slop patterns -- it is dense, operationally correct, and avoids gratuitous documentation. The _FrozenLinearCrossEntropy autograd function correctly discards weight gradients and properly promotes to float32 for gradient scaling. The fused AdamW gradient clipping via grad_scale uses PyTorch's documented interface. The SQLite coordinator schema is well-designed with proper foreign keys and indices. DCO sign-off is present on all commits.
No line-level findings survived filtering (all below confidence 80 after verification). All 12 CodeRabbit findings were addressed in follow-up commits. The documentation corrections (request-timeout default 180->120, max-anchors 256->3072) are accurate -- they match the actual code on main.
Recommend splitting per the author's offer before detailed maintainer review.
🤖 Generated with Claude Code using the /pr-review skill
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
This PR introduces a well-designed bounded asynchronous fan-out system for independent consumer trainers. The SQLite coordinator, filesystem artifact cache, and windowed-cursor semantics are sound, and the existing CodeRabbit findings have been addressed. The primary concern is scope.
Splittability. At ~11.5k additions across 37 files, this PR bundles at least four independently reviewable units: (1) the artifact cache and windowed coordinator, (2) training integration (data.py, dataloader.py, trainer.py), (3) DFlash consumer optimizations (fused CE, optimizer backend, metrics), and (4) the benchmark harness with GPU monitoring. The author has already offered to split into a dependency-ordered series. I would recommend taking that offer -- smaller PRs get tighter review and merge faster, and each unit here can be validated independently.
Issue assignment. Per CONTRIBUTING.md, significant changes (this PR modifies the CLI/API, adds new modules, and touches 37 files) should reference an assigned issue with a maintainer green-light. Issue #810 exists and was opened by the author, but has no assignees. A maintainer should explicitly assign the issue to signal alignment before this merges.
Two minor inline findings below. Neither is a blocker.
🤖 Generated with Claude Code using the /pr-review skill
| def write_report(report: dict[str, Any], path: Path) -> None: | ||
| path.parent.mkdir(parents=True, exist_ok=True) | ||
| temporary = path.with_suffix(path.suffix + ".tmp") | ||
| temporary.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") | ||
| temporary.replace(path) |
There was a problem hiding this comment.
write_report uses a temporary-then-rename pattern without flush()/fsync(), while _atomic_write_json in the sibling gpu_monitor.py module correctly calls both before the rename. On a crash, the report file could be zero-length or truncated despite the rename having completed from the kernel’s perspective. Since this is the primary benchmark output, it should match the durability guarantees used elsewhere in this package:
def write_report(report: dict[str, Any], path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
with temporary.open("w") as output:
json.dump(report, output, indent=2, sort_keys=True)
output.write("\n")
output.flush()
os.fsync(output.fileno())
temporary.replace(path)| step_ms: float | ||
|
|
||
|
|
||
| def _percentile(values: list[float], percentile: float) -> float: |
There was a problem hiding this comment.
This _percentile is nearly identical to the one in gpu_monitor.py (line 429). Both live in the same speculators.benchmarks package. Extract a single shared utility to avoid the DRY violation.
Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn>
|
Follow-up
Focused validation passes: The implementation scope remains frozen. I will keep the branch maintained and wait for maintainer guidance on issue assignment and whether splitting the existing end-to-end change would materially help review; I will not restructure the PR or add Mooncake/new feature scope speculatively. |
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: yangtianjian <143272889+heiheiha798@users.noreply.github.com>
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
|
This pull request has merge conflicts that must be resolved before it can be |
Purpose
Resolves #810.
Add opt-in bounded asynchronous hidden-state fan-out for independent consumers:
steady-state windows in a reproducible benchmark harness;
monitor; and
version-pinned Liger fused linear cross-entropy backend.
All new training behavior is opt-in. Existing online/offline data paths and optimizer
defaults remain unchanged. The first shared-artifact backend requires POSIX
flock,same-filesystem atomic rename, and directory
fsync; it does not claim arbitrary NFSor GPU-direct support.
Scope and related Mooncake work
This PR does not add a Mooncake payload backend. Its shared-artifact cache stores
safetensors on a shared POSIX filesystem, and all 1P3C throughput, memory, and NVML
results below measure that filesystem implementation.
A related Mooncake-backed bounded asynchronous producer/consumer implementation is
publicly available in
sgl-project/SpecForge#707. It includes
zero-copy transport and fan-out lifetime ownership. That implementation provides useful
related validation for the lifecycle invariants and benchmark design here, but its
Mooncake code is not included in this PR.
Upstream Mooncake work already exists in #605 and #710, with #735 providing the merged
hidden-state transfer abstraction and current file backend. This PR does not claim first
Mooncake support or attempt to duplicate those open changes.
The coordinator in this PR currently carries concrete filesystem paths in its read and
eviction leases. Supporting Mooncake for the shared fan-out payload therefore requires a
follow-up refactor to opaque artifact references and a store protocol; it is not an
already supported configuration or a drop-in adapter today. That follow-up should be
coordinated with #605/#710.
Three aligned Qwen3-8B DFlash 1P3C runs (3072 tokens, 256 anchors, heterogeneous block
sizes 4/8/16, 30 warmup plus 150 measured steps per consumer) produced aggregate steady
throughput of 6.3735, 6.2856, and 6.5702 steps/s. Median per-consumer rates were
2.2643/2.3135/1.6978 steps/s, and producer utilization was 100% over the common steady
interval.
Tests
make qualityon Python 3.10, 3.12, and 3.13: passed (Ruff, formatting,Markdown, and Mypy).
181 passed, 1 skipped.1 passed.liger-kernel: installed-package import,lazy fused-CE import, training CLI help, benchmark config validation, and dependency
consistency all passed.
git diff --check, DCO validation, and sensitive/local-path scan: passed.An additional local
mkdocs build --strictgenerated the complete site but did notexit after generation. The same behavior and the same 17 pre-existing Griffe warnings
were reproduced at the upstream base commit; no warning points to a changed file.
Checklist
I have filled in:
Follow-up validation at
5dfac242d7156d,47c2805, and5dfac24.make qualitypasses in the Python 3.12 project environment: Ruff, formatting, Markdown, and mypy (183files).198 passed(19 warnings).THIRD_PARTY_NOTICES.