ulysses sp draft - #853
Closed
Roderick-Wu wants to merge 34 commits into
Closed
Conversation
Contributor
|
Important Review skippedIgnore keyword(s) in the title. ⛔ Ignored keywords (3)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
|
The quality checks have failed. Please run |
Roderick-Wu
force-pushed
the
feat/ulysses-sp
branch
2 times, most recently
from
July 23, 2026 20:38
3b01abb to
62b0b5e
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/speculators/train/trainer.py (1)
294-306: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winGate or fix SP with the DDP path.
With
sp_size > 1, each rank trains on a sequence-chunked sample, but_setup_model_ddp()never wiresregister_sp_gradient_hooks(self.model). DDP with noprocess_groupuses the default group and averages gradients acrossdp_size * sp_sizeranks, so the optimizer receives gradients scaled down bysp_size. Add CLI validation to require--fsdp-shardwhen--sp-size > 1, or support SP on DDP viaprocess_group=get_dp_group()plus SP gradient hooks.🤖 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 `@src/speculators/train/trainer.py` around lines 294 - 306, Ensure sequence parallelism is not used with the DDP setup: add CLI/config validation requiring fsdp_shard whenever sp_size is greater than 1, before model setup proceeds. Keep the existing _setup_model_ddp and _setup_model_fsdp paths unchanged otherwise.Source: Path instructions
🧹 Nitpick comments (1)
src/speculators/models/metrics.py (1)
507-557: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider adding a numerical test for SP gradient correctness.
The
dist.all_reduce-then-local-backward()trick used here (non-differentiable collective on the numerator/denominator, relying on the local division backward using the post-all-reduce value as the correct chain-rule coefficient, combined withregister_sp_gradient_hooks's subsequent gradient sum) is mathematically sound but subtle and easy to silently regress. A small test comparing gradients from a 2-rank SP-split forward/backward against a single-rank full-sequence reference would give strong confidence this training-critical path stays correct across future refactors.As per path instructions for
src/speculators/train/**/*.pyto verify "loss computation handles padding masks properly," this loss/gradient path is exactly the kind of distributed training correctness surface that benefits from explicit test coverage.🤖 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 `@src/speculators/models/metrics.py` around lines 507 - 557, Add a focused numerical distributed test for loss_function with a 2-rank SP-split sequence, comparing the resulting gradients after backward and SP gradient synchronization against a single-rank full-sequence reference. Cover padding through loss_mask and exercise the all_reduce path when get_sp_size() > 1, preserving the expected masked loss and gradient equivalence.Source: Path instructions
🤖 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/train.py`:
- Around line 526-544: Add an early validation in the sequence-parallel setup
around maybe_setup_distributed and the existing args.sp_size checks: when
args.sp_size > 1, require is_distributed() to be true and raise a clear
ValueError otherwise, matching the existing --fsdp-shard guard. Keep the
divisibility, supported speculator type, and DFlash anchor validations unchanged
for valid distributed launches.
- Around line 939-950: Update the help text for the --sp-size argument in the
parser configuration to state that sequence parallelism is supported for both
eagle3 and dflash, matching the allowed speculator types in main().
In `@src/speculators/models/attention.py`:
- Around line 49-67: Before the query path in the sequence-parallel block,
validate that the query head count divides evenly by sp_size and apply the same
head replication handling used for K/V when it does not. Update the flow around
maybe_replicate_kv_heads and ulysses_scatter so query is safe to scatter across
all SP ranks, while preserving the existing K/V behavior.
In `@src/speculators/models/eagle3/core.py`:
- Around line 203-211: Update Eagle3’s loss-return path to divide scalar metrics
whose keys end in “_total” by sp_size before returning them, matching DFlash’s
sequence-parallel normalization. Apply this to metrics such as
metrics["loss_total"] while leaving non-total metrics and the existing
hidden_states/document_ids handling unchanged.
In `@src/speculators/train/sequence_parallel.py`:
- Around line 98-119: Track the tensor that first determines full_seq_len in the
sequence-parallel batching logic, and use that tensor’s device when creating
out["full_seq_len"]. Update the relevant loop state and final tensor
construction so they no longer depend on the loop-scoped tensor variable, while
preserving existing handling for split and full-length keys.
---
Outside diff comments:
In `@src/speculators/train/trainer.py`:
- Around line 294-306: Ensure sequence parallelism is not used with the DDP
setup: add CLI/config validation requiring fsdp_shard whenever sp_size is
greater than 1, before model setup proceeds. Keep the existing _setup_model_ddp
and _setup_model_fsdp paths unchanged otherwise.
---
Nitpick comments:
In `@src/speculators/models/metrics.py`:
- Around line 507-557: Add a focused numerical distributed test for
loss_function with a 2-rank SP-split sequence, comparing the resulting gradients
after backward and SP gradient synchronization against a single-rank
full-sequence reference. Cover padding through loss_mask and exercise the
all_reduce path when get_sp_size() > 1, preserving the expected masked loss and
gradient equivalence.
🪄 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 Plus
Run ID: e349f1a9-7859-4fe2-9723-167fa4f43b0e
📒 Files selected for processing (13)
.gitignorescripts/train.pysrc/speculators/models/attention.pysrc/speculators/models/dflash/core.pysrc/speculators/models/dflash/model_definitions.pysrc/speculators/models/dspark/core.pysrc/speculators/models/eagle3/core.pysrc/speculators/models/metrics.pysrc/speculators/models/mtp/core.pysrc/speculators/models/peagle/core.pysrc/speculators/train/distributed.pysrc/speculators/train/sequence_parallel.pysrc/speculators/train/trainer.py
Comment on lines
+49
to
+67
| from speculators.train.distributed import get_sp_group, get_sp_size # noqa: PLC0415 | ||
| from speculators.train.sequence_parallel import ( # noqa: PLC0415 | ||
| maybe_replicate_kv_heads, | ||
| ulysses_gather, | ||
| ulysses_scatter, | ||
| ) | ||
|
|
||
| sp_size = get_sp_size() | ||
| use_sp = sp_size > 1 | ||
|
|
||
| if use_sp: | ||
| sp_group = get_sp_group() | ||
| assert sp_group is not None, "SP group did not initialize for sequence parallelism, something went wrong!" | ||
|
|
||
| key, value = maybe_replicate_kv_heads(key, value, sp_size) | ||
| query = ulysses_scatter(query, sp_group, sp_size) | ||
| key = ulysses_scatter(key, sp_group, sp_size) | ||
| value = ulysses_scatter(value, sp_group, sp_size) | ||
|
|
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does torch.chunk raise an error or silently produce uneven chunks when a dimension isn't evenly divisible?
💡 Result:
No, torch.chunk does not raise an error when a dimension is not evenly divisible by the specified number of chunks [1][2]. Instead, it handles this case by producing chunks of the same size, except for the last one, which may be smaller [1][3][2]. Importantly, the function may return fewer than the specified number of chunks if such a division is not possible given the requested parameters [1][2]. The official PyTorch documentation explicitly warns that this function "may return fewer than the specified number of chunks" and suggests using torch.tensor_split if you require exactly the specified number of chunks [1][2]. The underlying implementation performs an integer division calculation: split_size = (tensor.size(dim) + chunks - 1) // chunks [4]. This logic results in a number of chunks that can be less than requested depending on the input size and the value of chunks [4].
Citations:
- 1: https://docs.pytorch.org/docs/stable/generated/torch.chunk.md
- 2: https://docs.pytorch.org/docs/2.6/generated/torch.chunk.html
- 3: https://docs.pytorch.org/docs/2.3/generated/torch.chunk.html
- 4: https://github.com/pytorch/pytorch/blob/32b38415535a7a0311cb644b3fddad47b438c401/torch/functional.py
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate and inspect the relevant implementation and call sites.
printf 'Files:\n'
git ls-files | rg '(^|/)(attention|sequence_parallel)\.py$|distributed' || true
printf '\nattention.py outline:\n'
ast-grep outline src/speculators/models/attention.py --view expanded || true
printf '\nsequence_parallel.py outline:\n'
ast-grep outline src/speculators/train/sequence_parallel.py --view expanded || true
printf '\nRelevant attention.py lines 1-120:\n'
sed -n '1,120p' src/speculators/models/attention.py
printf '\nRelevant sequence_parallel.py lines 1-220:\n'
sed -n '1,220p' src/speculators/train/sequence_parallel.py
printf '\nSearch for num_query_heads / query heads validation in sequence_parallel:\n'
rg -n "num_query_heads|query_heads|chunk\\(|scatter_dim=1|maybe_replicate_kv_heads|AllToAllSP" src/speculators/train src/speculators/models -SRepository: vllm-project/speculators
Length of output: 14438
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import math
num_heads_vals = [3, 5, 7, 8, 9, 12, 13]
sp_vals = [2, 3, 5]
for h in num_heads_vals:
for sp in sp_vals:
chunks = [1 + (h - 1 - i) // sp for i in range(sp)]
print(f"h={h} sp={sp}: splits=[{', '.join(f'{x}:{h*i//sp + x}' for x in chunks)}], last_is_smaller={chunks[-1] < chunks[0]}")
PYRepository: vllm-project/speculators
Length of output: 330
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
num_heads_vals = [3, 5, 7, 8, 9, 12, 13]
sp_vals = [2, 3, 5]
for h in num_heads_vals:
for sp in sp_vals:
chunks = []
size = h
for i in range(sp):
split_size = (size + sp - 1) // sp
chunk_size = min(split_size, size)
chunks.append(chunk_size)
size -= split_size
print(f"h={h} sp={sp}: chunks={chunks}, first={chunks[0] if chunks else 0}, last={chunks[-1]}, last_smaller={chunks[-1] < chunks[0]}")
PYRepository: vllm-project/speculators
Length of output: 1514
Validate that num_query_heads divides cleanly by sp_size before scattering.
_AllToAllSP.forward uses input_tensor.chunk(world_size, dim=scatter_dim), and torch.chunk creates uneven chunks when the dimension is not evenly divisible. maybe_replicate_kv_heads protects the K/V path, but query is scattered directly by ulysses_scatter, so mismatched GQA query head counts can silently give some SP ranks fewer query heads and produce broken attention results. Add an explicit divide/replication check before scattering query, matching the existing KV-head path.
🤖 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 `@src/speculators/models/attention.py` around lines 49 - 67, Before the query
path in the sequence-parallel block, validate that the query head count divides
evenly by sp_size and apply the same head replication handling used for K/V when
it does not. Update the flow around maybe_replicate_kv_heads and ulysses_scatter
so query is safe to scatter across all SP ranks, while preserving the existing
K/V behavior.
Source: Path instructions
Comment on lines
+203
to
+211
| # With Ulysses SP, hidden_states is split but document_ids is full. | ||
| local_seq_len = hidden_states.shape[1] | ||
| total_seq_len = document_ids.shape[1] | ||
|
|
||
| if position_ids is None: | ||
| position_ids = 1 + torch.arange( | ||
| total_seq_len, dtype=torch.long, device=device | ||
| local_seq_len, dtype=torch.long, device=device | ||
| ).unsqueeze(0) | ||
| # shape: [1, total_seq_len] | ||
| # shape: [1, local_seq_len] |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate relevant files =="
git ls-files | rg '(^|/)src/speculators/models/(eagle3|dflash)/core\.py$|(^|/)trainer\.py$|scripts/train.py$|.*compute_metrics.*|.*loss_function.*' || true
echo
echo "== eagle3 core outline =="
ast-grep outline src/speculators/models/eagle3/core.py --view compact || true
echo
echo "== eagle3 core relevant lines =="
cat -n src/speculators/models/eagle3/core.py | sed -n '160,250p'
echo
echo "== dflash core relevant total metric handling =="
if [ -f src/speculators/models/dflash/core.py ]; then
cat -n src/speculators/models/dflash/core.py | sed -n '1,260p' | rg -n "sp_size|_total|metrics|rank|dist|reduce|loss" -C 3 || true
fi
echo
echo "== trainer metric reduction/normalization =="
if [ -f src/speculators/trainers/trainer.py ]; then
cat -n src/speculators/trainers/trainer.py | sed -n '1,260p' | rg -n "loss_total|loss_sum|metrics|reduce|normalize_counted_metrics|sp_split|Ulysses|Sequence Parallel|dist|SUM" -C 4 || true
fi
echo
echo "== search total/loss metrics in codebase =="
rg -n 'loss_total|loss_sum|normalize_counted_metrics|sp_size|sp_splits_batch|document_ids|position_ids|metrics\["loss' -S .Repository: vllm-project/speculators
Length of output: 36916
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== deterministic call/reduction probe: parse total metric definitions and loss function args =="
python3 - <<'PY'
from pathlib import Path
import re
files = [
Path("src/speculators/models/eagle3/core.py"),
Path("src/speculators/models/dflash/core.py"),
Path("src/speculators/trainers/trainer.py"),
]
for p in files:
print(f"\n-- {p} --")
s = p.read_text()
lines = s.splitlines()
for i,l in enumerate(lines,1):
if "metrics[" in l and ("loss_total" in l or "loss_sum" in l or "_total" in l):
context = lines[max(1,i-5):i+6]
print(f"lines {max(1,i-5)}-{i+5}:")
for j,ctx in enumerate(context, max(1,i-5)):
print(f"{j}: {ctx}")
if "def compute_metrics" in l or "loss_function" in l or "normalize_counted_metrics" in l or "dist.reduce" in l:
print(f"metric/dist line {i}: {l}")
# Extract compute_metrics signatures where present using simple line scanning
for p in files:
s=p.read_text()
for m in re.finditer(r'(?m)^(\s*def compute_metrics\s*\([^)]*\)).*:(?ms).*?$|(?m)^\s*def compute_metrics\b', s):
start=m.start()+1
# naive next def/class/EOF
txt=s[m.start():]
end=len(txt)
for nextpat in (r'(?m)^(\s*def |\s*class )'):
mm=re.search(nextpat, txt[1:])
if mm and mm.start()+1 < end: end=mm.start()+1+1
print(f"\ncompute_metrics in {p}:")
print(txt[:end].splitlines()[:12])
PYRepository: vllm-project/speculators
Length of output: 2053
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== trainer compute/normalize path =="
cat -n src/speculators/train/trainer.py | sed -n '440,500p'
cat -n src/speculators/train/trainer.py | sed -n '535,580p'
echo
echo "== utils normalize_counted_metrics =="
cat -n src/speculators/train/utils.py | sed -n '1,120p'
echo
echo "== exact trainer metric reduction searches =="
rg -n "forward\\(|compute\\(|normalize_counted_metrics|reduce\\(|metrics|loss_total|loss_sum|sp_size" src/speculators/train/trainer.py -C 4
echo
echo "== dflash forward tail with SP total normalization =="
cat -n src/speculators/models/dflash/core.py | sed -n '480,528p'
echo
echo "== scripts train SP restrictions =="
cat -n scripts/train.py | sed -n '520,545p'Repository: vllm-project/speculators
Length of output: 19482
Normalize Eagle3 _total metrics under sequence parallelism.
trainer.py reduces every metric across ranks with ReduceOp.SUM before normalize_counted_metrics(), and split_batch_for_sp() gives Eagle3 each rank a split hidden-state batch while keeping full document_ids. Since Eagle3 still emits metrics["loss_total"] = torch.tensor(1.0, ...) from every SP rank, the aggregated loss ratio will be divided by a total count that is inflated by sp_size; mirror DFlash’s sp_size normalization for _total scalar metrics in the same loss-return path.
🤖 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 `@src/speculators/models/eagle3/core.py` around lines 203 - 211, Update
Eagle3’s loss-return path to divide scalar metrics whose keys end in “_total” by
sp_size before returning them, matching DFlash’s sequence-parallel
normalization. Apply this to metrics such as metrics["loss_total"] while leaving
non-total metrics and the existing hidden_states/document_ids handling
unchanged.
Roderick-Wu
marked this pull request as ready for review
July 27, 2026 03:26
Roderick-Wu
requested review from
fynnsu,
orestis-z,
rahul-tuli and
shanjiaz
as code owners
July 27, 2026 03:26
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
Design-level review
This PR adds Ulysses sequence parallelism across Eagle3, DFlash, DSpark, and PEagle. The core approach is sound: a torch.library.custom_op wrapping dist.all_to_all with proper autograd and torch.compile support, combined with per-model strategies for splitting work across SP ranks (batch-level splitting for Eagle3 via the trainer, model-internal anchor partitioning for DFlash/PEagle). The FSDP integration correctly scopes sharding to the DP sub-mesh and registers post-accumulate gradient hooks for SP gradient reduction.
However, there are several process and coverage concerns before this is reviewable in earnest:
Missing motivation and PR description
The title says "ulysses sp draft" and the Purpose / Tests / Checklist sections are empty. For a change that touches 14 files including core training infrastructure (distributed.py, trainer.py), all model forward passes, and loss computation, the PR description needs to explain:
- Why this is needed (long-context training? memory reduction benchmarks?).
- How it was validated (what models, what
sp_size, what seq lengths). - How it relates to the existing SP Ulysses PR #640 (by a maintainer), which is still open and tracked on the Q3 roadmap.
Per CONTRIBUTING.md, changes affecting 3+ files should reference an assigned issue. No issue is linked here.
No test coverage
The PR adds 294 lines of new code in sequence_parallel.py, modifies 5 model forward passes, changes distributed gradient handling, and adds a new attention function -- all without a single test. The SP gradient path in particular (non-differentiable all_reduce on numerator/denominator in loss_function, combined with register_sp_gradient_hooks for parameter gradients) is subtle and easy to silently regress. At minimum:
- A unit test for
split_batch_for_spcovering key splitting, full-key preservation, and thefull_seq_lenmetadata. - A unit test for
maybe_replicate_kv_headscovering the divisibility checks and replication logic. - A 2-rank numerical test comparing SP-split forward/backward against a single-rank full-sequence reference to verify gradient equivalence.
Bundled unrelated changes
The PR mixes the SP feature with unrelated additions (see inline comments). These should be separate PRs or removed.
🤖 Generated with Claude Code using the /pr-review skill
| @@ -0,0 +1,74 @@ | |||
| """Wrap vllm serve with faulthandler to catch native crashes. | |||
There was a problem hiding this comment.
This 74-line debugging helper (with hardcoded defaults, , ) adds maintenance surface without clear value to the repo. If it is useful for the SP feature, it should be in a separate PR with a description explaining the use case. Otherwise, consider keeping it as a local script outside the repo.
Comment on lines
1357
to
1360
| print(f"Peak GPU memory: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB") | ||
|
|
||
|
|
||
| # RUN WITH: |
There was a problem hiding this comment.
This unconditional print at the end of the __main__ block looks like a debugging artifact. It will fire on every training run, including production runs and CI. If this is useful, it should be gated behind a flag (e.g. --log-peak-memory) or moved behind the existing log_freq mechanism rather than raw print.
|
Link Check Results (DOCS) Summary
Errors per inputErrors in docs/index.md
Redirects per inputRedirects in docs/community/contact_us.mdRedirects in docs/community/index.mdRedirects in docs/index.md
Last updated: 982377b |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@src/speculators/models/peagle/core.py`:
- Around line 56-73: Update _broadcast_and_partition_anchors and _SPPartition to
preserve the pre-padding sample count or an equivalent validity mask alongside
padded anchor/depth data. Use that validity information to exclude padded nodes
from attention and prevent padded local_orig_positions from contributing to loss
or metrics. Apply the validity mask to sampled_loss_mask before loss and
denominator calculations, while preserving existing behavior for original
samples.
In `@src/speculators/train/config/schema.py`:
- Around line 404-411: Update the schema validation for DSpark configurations so
that when DSpark uses sequence parallelism greater than one, total_seq_len must
be divisible by sp_size multiplied by block_size. Preserve existing validation
for other models and for non-SP DSpark runs, and emit a clear validation error
when the new alignment requirement is violated.
In `@src/speculators/train/sequence_parallel.py`:
- Around line 46-51: Validate in all_to_all_sp and ulysses_scatter that
input_tensor.shape[scatter_dim] is divisible by sp_size before calling
torch.chunk, raising a clear error for incompatible dimensions. Also add
equivalent validation during launch setup so incompatible query-head and sp_size
combinations are rejected before execution.
🪄 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 Plus
Run ID: ddbffc58-14fe-4e59-b098-946e4ef3a3ef
📒 Files selected for processing (9)
scripts/train.pysrc/speculators/models/attention.pysrc/speculators/models/dflash/core.pysrc/speculators/models/dspark/metrics.pysrc/speculators/models/metrics.pysrc/speculators/models/peagle/core.pysrc/speculators/train/config/schema.pysrc/speculators/train/sequence_parallel.pysrc/speculators/train/trainer.py
💤 Files with no reviewable changes (1)
- src/speculators/models/dflash/core.py
🚧 Files skipped from review as they are similar to previous changes (4)
- src/speculators/models/metrics.py
- src/speculators/train/trainer.py
- src/speculators/models/attention.py
- scripts/train.py
Comment on lines
+56
to
+73
| remainder = total_sampled % sp_size | ||
| if remainder != 0: | ||
| pad_len: int = sp_size - remainder | ||
| anchor_pos = torch.nn.functional.pad(anchor_pos, (0, pad_len)) | ||
| depth = torch.nn.functional.pad(depth, (0, pad_len)) | ||
|
|
||
| total_sampled = anchor_pos.shape[0] | ||
| n_per_rank = total_sampled // sp_size | ||
| local_start = sp_rank * n_per_rank | ||
| return _SPPartition( | ||
| full_anchor_pos=anchor_pos, | ||
| full_depth=depth, | ||
| total_sampled=total_sampled, | ||
| local_anchor_pos=anchor_pos[local_start : local_start + n_per_rank], | ||
| local_depth=depth[local_start : local_start + n_per_rank], | ||
| local_orig_positions=anchor_pos[local_start : local_start + n_per_rank] | ||
| + depth[local_start : local_start + n_per_rank], | ||
| local_sampled=n_per_rank, |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate file =="
fd -a 'core.py' . | sed 's#^\./##' | grep 'src/speculators/models/peagle/core.py' || true
echo "== outline =="
ast-grep outline src/speculators/models/peagle/core.py --view expanded || true
echo "== relevant sections =="
nl -ba src/speculators/models/peagle/core.py | sed -n '1,260p'
echo "== usages of _SPPartition / broadcast function / metrics =="
rg -n "_SPPartition|_broadcast_and_partition_anchors|compute_metrics|anchor_pos|local_orig_positions|full_anchor_pos|local_depth|full_depth" src/speculators/models/peagle src -g '*.py' || trueRepository: vllm-project/speculators
Length of output: 1068
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== relevant core.py sections =="
awk '{printf "%6d %s\n", NR, $0}' src/speculators/models/peagle/core.py | sed -n '1,220p'
echo "== remaining core.py =="
awk '{printf "%6d %s\n", NR, $0}' src/speculators/models/peagle/core.py | sed -n '220,420p'
echo "== usages in speculators models =="
rg -n "_SPPartition|_broadcast_and_partition_anchors|forward\\(|mask|attention|compute_metrics|local_orig_positions|full_anchor_pos|full_depth|local_anchor_pos|local_depth|local_sampled|total_sampled" src/speculators/models/peagle/src/speculators src/speculators/models src/speculators -g '*.py' || true
echo "== file size and imports =="
wc -l src/speculators/models/peagle/core.py
sed -n '1,70p' src/speculators/models/peagle/core.pyRepository: vllm-project/speculators
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== relevant core.py sections =="
awk '{printf "%6d %s\n", NR, $0}' src/speculators/models/peagle/core.py | sed -n '1,220p'
echo "== remaining core.py =="
awk '{printf "%6d %s\n", NR, $0}' src/speculators/models/peagle/core.py | sed -n '220,420p'
echo "== usages in speculators models =="
rg -n "_SPPartition|_broadcast_and_partition_anchors|forward\\(|mask|attention|compute_metrics|local_orig_positions|full_anchor_pos|full_depth|local_anchor_pos|local_depth|local_sampled|total_sampled" src/speculators/models/peagle src/speculators -g '*.py' || true
echo "== file size and imports =="
wc -l src/speculators/models/peagle/core.py
sed -n '1,70p' src/speculators/models/peagle/core.pyRepository: vllm-project/speculators
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== peagle attention =="
awk '{printf "%6d %s\n", NR, $0}' src/speculators/models/peagle/attention.py | sed -n '1,120p'
echo "== peagle metrics =="
awk '{printf "%6d %s\n", NR, $0}' src/speculators/models/peagle/metrics.py | sed -n '1,140p'
echo "== shared metric loss functions relevant =="
awk '{printf "%6d %s\n", NR, $0}' src/speculators/models/metrics.py | sed -n '460,555p'
echo "== focused code scan for _SPPartition/padded validity =="
python3 - <<'PY'
from pathlib import Path
import ast
for path in [Path("src/speculators/models/peagle/core.py"), Path("src/speculators/models/peagle/attention.py"), Path("src/speculators/models/peagle/metrics.py")]:
print(f"\n--- {path} ---")
with open(path) as f:
src = f.read()
tree = ast.parse(src, filename=str(path))
for node in ast.walk(tree):
if isinstance(node, ast.NamedExpr):
print(f"walrus at line {node.lineno}: {ast.unparse(node.target)}")
elif isinstance(node, ast.Name) and node.id in {"valid", "anchor_valid", "padding", "pad", "mask", "valid_mask"}:
print(f"name {node.id} at line {node.lineno}")
with open("src/speculators/models/peagle/core.py") as f:
lines = f.readlines()
start = _broadcast=line for line in lines[:75]
if False: pass
for i, line in enumerate(lines, 1):
if "valid" in line.lower() or "valid_mask" in line or "anchor_valid" in line or "padding_mask" in line or "is_mask" in line or "pad" in line:
print(f"{i}: {line.rstrip()}")
PYRepository: vllm-project/speculators
Length of output: 11055
Carry and apply the original sample count for parallel COD samples.
_broadcast_and_partition_anchors pads non-divisible anchor/depth arrays, but _SPPartition only exposes total_sampled; the full mask still treats padded entries as position-0, padded rank-local orig_positions read from index 0 and can inject duplicated loss_mask[0] into loss/metrics denominators, and padded attention entries can participate via document-mask checks. Keep an original-count or validity mask through _SPPartition, exclude padded nodes in attention, and mask them out of sampled_loss_mask before computing loss/metrics.
🤖 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 `@src/speculators/models/peagle/core.py` around lines 56 - 73, Update
_broadcast_and_partition_anchors and _SPPartition to preserve the pre-padding
sample count or an equivalent validity mask alongside padded anchor/depth data.
Use that validity information to exclude padded nodes from attention and prevent
padded local_orig_positions from contributing to loss or metrics. Apply the
validity mask to sampled_loss_mask before loss and denominator calculations,
while preserving existing behavior for original samples.
Source: Path instructions
|
The quality checks have failed. Please run |
Roderick-Wu
commented
Jul 27, 2026
Comment on lines
-110
to
146
| bsz, q_len = hidden_states.shape[:-1] | ||
| ctx_len = target_hidden.shape[1] | ||
| q = self.q_proj(hidden_states) | ||
| q = q.view(bsz, q_len, -1, self.head_dim) | ||
| q = q.unflatten(-1, (-1, self.head_dim)) | ||
| q = self.q_norm(q).transpose(1, 2) | ||
| # This is the main difference from the usual attention mechanism. | ||
| k_ctx = self.k_proj(target_hidden) | ||
| k_noise = self.k_proj(hidden_states) | ||
| v_ctx = self.v_proj(target_hidden) | ||
| v_noise = self.v_proj(hidden_states) | ||
| k = torch.cat([k_ctx, k_noise], dim=1).view( | ||
| bsz, ctx_len + q_len, -1, self.head_dim | ||
| ) | ||
| # note the length becomes context length + block size | ||
| v = torch.cat([v_ctx, v_noise], dim=1).view( | ||
| bsz, ctx_len + q_len, -1, self.head_dim | ||
| ) | ||
| k = torch.cat([k_ctx, k_noise], dim=1).unflatten(-1, (-1, self.head_dim)) | ||
| v = torch.cat([v_ctx, v_noise], dim=1).unflatten(-1, (-1, self.head_dim)) | ||
| k = self.k_norm(k).transpose(1, 2) | ||
| v = v.transpose(1, 2) | ||
| cos, sin = position_embeddings | ||
| q, k = apply_rotary_pos_emb(q, k, cos, sin) | ||
| if past_key_values is not None: | ||
| cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} | ||
| k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) | ||
| attn_fn: Callable = eager_attention_forward | ||
| if ( | ||
| self.config._attn_implementation is not None # noqa: SLF001 | ||
| and self.config._attn_implementation != "eager" # noqa: SLF001 | ||
| ): | ||
| attn_fn = ALL_ATTENTION_FUNCTIONS[ | ||
| self.config._attn_implementation # noqa: SLF001 | ||
| ] | ||
| attn_output, attn_weights = attn_fn( | ||
| self, | ||
| q, | ||
| k, | ||
| v, |
Collaborator
Author
There was a problem hiding this comment.
These shape reshuffles are by claude and I didn't look too closely into
Roderick-Wu
commented
Jul 27, 2026
Comment on lines
139
to
372
| @@ -293,7 +299,7 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None: | |||
| if loaded_hs is None: | |||
| raise ValueError(f"Failed to load hidden states for handle {handle}") | |||
|
|
|||
| check_hidden_states(loaded_hs, dataset_item["input_ids"].tolist()) | |||
| check_hidden_states(loaded_hs, client_item["input_ids"]) | |||
|
|
|||
| file_idx = self._map_to_file_idx(index) | |||
| match self.on_generate: | |||
| @@ -341,23 +347,28 @@ def _get_raw_data(self, index): | |||
| # "token_ids": [seq_len] | |||
| # } | |||
|
|
|||
| if not torch.equal(loaded_hs["token_ids"], self.data[index]["input_ids"]): | |||
| cached_ids = loaded_hs["token_ids"] | |||
| full_ids = self.data[index]["input_ids"] | |||
| # Cached hidden states may cover N-1 tokens (last token trimmed | |||
| # before sending to vLLM); accept any valid prefix of the full ids. | |||
| if not torch.equal(cached_ids, full_ids[: len(cached_ids)]): | |||
| warnings.warn( | |||
| f"Loaded token ids {loaded_hs['token_ids']} for index {index} don't" | |||
| f"match input ids {self.data[index]['input_ids']}", | |||
| f"Loaded token ids {cached_ids} for index {index} don't" | |||
| f"match input ids {full_ids}", | |||
| stacklevel=1, | |||
| ) | |||
| return None | |||
|
|
|||
| hs_len = cached_ids.shape[0] | |||
| return { | |||
| "hidden_states": loaded_hs["hidden_states"][:, :-1].flatten( | |||
| 1 | |||
| ), # [seq_len, 3 * hidden_size] | |||
| "input_ids": loaded_hs["token_ids"], # [seq_len] | |||
| ), # [hs_len, 3 * hidden_size] | |||
| "input_ids": cached_ids, # [hs_len] | |||
| "verifier_last_hidden_states": loaded_hs["hidden_states"][ | |||
| :, -1 | |||
| ], # [seq_len, hidden_size] | |||
| "loss_mask": self.data[index]["loss_mask"], # [seq_len] | |||
| ], # [hs_len, hidden_size] | |||
| "loss_mask": self.data[index]["loss_mask"][:hs_len], # [hs_len] | |||
| } | |||
Collaborator
Author
There was a problem hiding this comment.
going to open a seperate pr for these
added 2 commits
July 27, 2026 17:09
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-01.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
Adds .github/CODEOWNERS to define review ownership. GitHub uses last-match-wins: the most specific matching path decides reviewers, not a union of all matches. The file is ordered broad → specific accordingly. Hierarchy: 1. Package-wide defaults — `/src/speculators/`,`/tests/`, `/scripts/` owned by all core maintainers (@fynnsu @shanjiaz @rahul-tuli @orestis-z), acting as a catch-all. 2. Subsystem overrides — convert/, data_generation/, train/ narrow to whoever owns that subsystem. 3. Algorithm-specific overrides — convert/mtp/, convert/eagle/, models/mtp/, models/peagle/, models/eagle3/, models/dflash/ route to the individual method's owner(s), overriding both subsystem and package-wide rules. 4. Tests mirror source — `/tests/` follows the same pattern, with e2e/regression/ and e2e/smoke/ also adding @dsikka. 5. Infra/process files — `.buildkite/`, `.github/`, and root config/docs (README.md, pyproject.toml, etc.) owned separately by @dsikka plus relevant maintainers. Example: a PR touching src/speculators/models/mtp/layer.py matches three rules — `/src/speculators/` (line 4), `/src/speculators/models/` (line 9), and `/src/speculators/models/mtp/` (line 10) — but only the last, most specific match applies, so reviewers are @rahul-tuli @fynnsu, not the full set of four maintainers. Signed-off-by: Dipika <dipikasikka1@gmail.com> Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
Mergify previously allowed any collaborator with write access to satisfy the single-approval (or second two-reviews approval) requirement. Restrict required approvals to a fixed list of reviewers (fynnsu, shanjiaz, dsikka, rahul-tuli, orestis-z). --------- Signed-off-by: Dipika <dipikasikka1@gmail.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
## Purpose - Update readme with latest models and an updated what's new section - Fix docs homepage I have filled in: - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [ ] The test plan/results, such as providing test command and pasting the results. - [ ] (Optional) The necessary documentation update. - [ ] I (a human) have written or reviewed the code in this pr to the best of my ability. Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
## Purpose Add a Tool-Call Regeneration subsection to the response-regeneration tutorial (after Multi-GPU Configurations). It documents: - the per-model `--tool-call-parser` / `--reasoning-parser` flags for Qwen3, Gemma 4, and gpt-oss, with values taken from the `vllm-project/recipes` source files (`models/Qwen/Qwen3-32B.yaml`, `models/Google/gemma-4-E2B-it.yaml`, `OpenAI/GPT-OSS.md`); - that this is semi-on-policy: the target regenerates the tool-call tokens on-policy, but tools are not executed. The i-th cached tool result from the source data is spliced positionally after the target's i-th regenerated call; - the limitation that a turn with parallel tool calls is truncated to the first call. ## Tests Docs-only change. `mdformat --check docs/user_guide/tutorials/response_regeneration.md` passes. ## Checklist I have filled in: - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [x] The test plan/results, such as providing test command and pasting the results. - [x] (Optional) The necessary documentation update. - [x] I (a human) have written or reviewed the code in this pr to the best of my ability. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com> Co-authored-by: shanjiaz <zsjwpianpian@gmail.com> Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
Implements the config-file-first CLI RFC (#833): replaces the ~600-line hand-written `argparse` surface in `scripts/train.py` with a typed, layered configuration system in a new `src/speculators/train/config/` package built on `pydantic-settings`. This is **phase 1** of the migration. The schema is now the single source of truth for every `train.py` tunable, and `main()` bridges to the existing model layer through a single `argparse.Namespace(**cfg.flatten())` adapter (retired in a later phase). No new dependencies — pydantic, pydantic-settings, and PyYAML are already in the stack. > **Reviewing?** The four commits are each self-contained and pass their own tests — read commit-by-commit using the **Reviewer reading order** below. `scripts/train.py` is a net **−660 lines** of parser; `main()`'s only changes are the `flatten()` bridge and dropping a now-redundant `float16` runtime guard (the rejection lives in the schema validator), so it can be skimmed. - `--config PATH` — load a stage-shaped YAML config; precedence is `flag > yaml > default`. - `--dump-config` — print the fully-resolved config as a round-trippable `run.yaml` and exit. - Every run writes `run.yaml` + `train_command.txt` next to its checkpoints, so a run reloads via `--config run.yaml` without hand-reconstructing flags. Existing flag-based recipes are unchanged — every flag, env var, and default is preserved: ```bash scripts/train.py \ --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \ --data-path ./data --save-path ./ckpt \ --speculator-type eagle3 --epochs 3 --lr 1e-4 --total-seq-len 4096 ``` `--dump-config` emits only the values a non-default layer supplied (stage-shaped, clean of provenance): ```yaml train: trainer: save_path: ./ckpt epochs: 3 data: total_seq_len: 4096 data_path: ./data optimizer: lr: 0.0001 speculator_type: eagle3 verifier: verifier_name_or_path: meta-llama/Llama-3.1-8B-Instruct ``` That file reloads identically via `--config run.yaml`, and re-dumping it is **byte-stable** (emission follows the pinned `flatten()` order, independent of flag-vs-yaml source). The flag surface is generated from the schema — one typed field becomes one flag, in a `--help` grouped by concern (`general`, `verifier`, `draft`, `data`, `generation`, `loss`, `optimizer`, `scheduler`, `trainer`, `logging`, and the per-algorithm groups `dflash` / `dspark` / `peagle` / `mtp`). `--config` / `--dump-config` lead the surface; every other flag, choice, bool style, and default is preserved from the legacy parser. ```bash $ scripts/train.py --help ``` <details> <summary>Full generated <code>--help</code> output</summary> ```text usage: train.py [-h] [--config PATH] [--dump-config] [--speculator-type SPECULATOR_TYPE] [--dry-run] [--seed SEED] [--deterministic-cuda] [--verifier-name-or-path VERIFIER_NAME_OR_PATH] [--trust-remote-code] [--from-pretrained FROM_PRETRAINED] [--draft-config DRAFT_CONFIG] [--num-layers NUM_LAYERS] [--draft-arch {llama,qwen3}] [--draft-hidden-act DRAFT_HIDDEN_ACT] [--draft-mrope-full-head-hack | --no-draft-mrope-full-head-hack] [--target-layer-ids TARGET_LAYER_IDS [TARGET_LAYER_IDS ...]] [--token-freq-path TOKEN_FREQ_PATH] [--draft-vocab-size DRAFT_VOCAB_SIZE] [--d2t-path D2T_PATH] [--t2d-path T2D_PATH] [--mask-token-id MASK_TOKEN_ID] [--norm-before-residual | --no-norm-before-residual] [--embed-requires-grad | --no-embed-requires-grad] [--norm-before-fc | --no-norm-before-fc] [--fc-norm] [--norm-output | --no-norm-output] [--sliding-window SLIDING_WINDOW] [--full-attention-indices FULL_ATTENTION_INDICES [FULL_ATTENTION_INDICES ...]] [--sliding-window-non-causal] [--draft-attn-impl {simple_flex_attention,sdpa,eager}] [--data-path DATA_PATH] [--hidden-states-backend {file}] [--hidden-states-path HIDDEN_STATES_PATH] [--total-seq-len TOTAL_SEQ_LEN] [--train-data-ratio TRAIN_DATA_RATIO] [--noise-std NOISE_STD] [--legacy-data] [--hidden-states-dtype HIDDEN_STATES_DTYPE] [--num-workers NUM_WORKERS] [--prefetch-factor PREFETCH_FACTOR] [--max-anchors MAX_ANCHORS] [--vllm-endpoint VLLM_ENDPOINT] [--on-missing {generate,skip,warn,raise}] [--on-generate {cache,delete}] [--request-timeout REQUEST_TIMEOUT] [--max-retries MAX_RETRIES] [--loss-fn LOSS_FN] [--ttt-steps TTT_STEPS] [--ttt-step-loss-decay TTT_STEP_LOSS_DECAY] [--optimizer {adamw,muon}] [--lr LR] [--weight-decay WEIGHT_DECAY] [--muon-lr MUON_LR] [--muon-momentum MUON_MOMENTUM] [--muon-weight-decay MUON_WEIGHT_DECAY] [--muon-ns-steps MUON_NS_STEPS] [--muon-adjust-lr-fn {original,match_rms_adamw}] [--scheduler-type {linear,cosine,none}] [--scheduler-warmup-steps SCHEDULER_WARMUP_STEPS] [--scheduler-warmup-ratio SCHEDULER_WARMUP_RATIO] [--scheduler-total-steps SCHEDULER_TOTAL_STEPS] [--scheduler-num-cosine-cycles SCHEDULER_NUM_COSINE_CYCLES] [--epochs EPOCHS] [--checkpoint-freq CHECKPOINT_FREQ] [--save-best] [--no-resume-from-checkpoint] [--log-freq LOG_FREQ] [--save-path SAVE_PATH] [--fsdp-shard] [--logger LOGGER] [--log-dir LOG_DIR] [--run-name RUN_NAME] [--block-size BLOCK_SIZE] [--sample-from-anchor | --no-sample-from-anchor] [--dflash-decay-gamma DFLASH_DECAY_GAMMA] [--per-position-loss-weight {fixed-exp-decay,dpace}] [--dpace-alpha DPACE_ALPHA] [--markov-rank MARKOV_RANK] [--markov-head-type {vanilla,gated,rnn}] [--enable-confidence-head | --no-enable-confidence-head] [--confidence-head-with-markov | --no-confidence-head-with-markov] [--confidence-head-alpha CONFIDENCE_HEAD_ALPHA] [--num-depths NUM_DEPTHS] [--down-sample-ratio DOWN_SAMPLE_RATIO] [--down-sample-ratio-min DOWN_SAMPLE_RATIO_MIN] [--num-speculative-steps NUM_SPECULATIVE_STEPS] [--step-weight-beta STEP_WEIGHT_BETA] options: -h, --help show this help message and exit --config PATH YAML config file (stage-shaped: trainer keys under a top-level 'train:' key; a bare mapping is also accepted). CLI flags override it. --dump-config Print the fully-resolved config as a valid run.yaml to stdout and exit, so a working invocation can be scaffolded into a shareable config. general: --speculator-type SPECULATOR_TYPE Type of speculator model to train (eagle3, dflash, dspark, peagle, mtp). --dry-run Build the speculator, initialize weights, save a checkpoint to --save-path, then exit before training. Useful to validate the config and weights (e.g. in vLLM) before a full run. Can be combined with --draft- config or --from-pretrained. --seed SEED Random seed for reproducibility. --deterministic-cuda Set cuda to deterministic mode. This may impact performance. verifier: --verifier-name-or-path VERIFIER_NAME_OR_PATH Path or HF id of the verifier/target model. --trust-remote-code Allow executing code from HF Hub when loading the verifier's tokenizer. draft: --from-pretrained FROM_PRETRAINED Path or HF id of a pretrained draft. May also point to a local directory containing only a config.json, in which case a fresh draft is initialized from that full speculator config. Takes precedence over and is mutually exclusive with --draft-config and the decoder-shaping flags (--num-layers, --draft-arch, --draft-hidden-act, --sliding-window, --full- attention-indices). --draft-config DRAFT_CONFIG HF id, directory, or JSON path of a decoder config (LlamaConfig for eagle3/peagle, Qwen3Config for dflash) to use as the draft transformer_layer_config; the rest of the speculator is built from the other CLI args. Mutually exclusive with --from-pretrained and the decoder-shaping flags (--num-layers, --draft-arch, --draft-hidden-act, --sliding-window, --full- attention-indices). --num-layers NUM_LAYERS Number of draft decoder layers to synthesize. --draft-arch {llama,qwen3} Architecture for draft decoder layers (default: 'llama' for eagle3, 'qwen3' otherwise). --draft-hidden-act DRAFT_HIDDEN_ACT Activation function for draft decoder layers. Defaults to 'silu'. Qwen3 layers of dflash expect 'silu' for vLLM deployment. Leave as None to fall back to the verifier's activation function. --draft-mrope-full-head-hack, --no-draft-mrope-full-head-hack For MRoPE configs with partial_rotary_factor < 1, rescale mrope_section and set partial_rotary_factor=1.0 so HF training and vLLM inference use equivalent full-head rotary semantics. --target-layer-ids TARGET_LAYER_IDS [TARGET_LAYER_IDS ...] (Optional) space-separated list of integer layer ids. Defaults to [2, n//2, n-3, n]. Must be set explicitly if custom values were used to launch vllm. --token-freq-path TOKEN_FREQ_PATH Path to token frequency distribution file (.pt). Used with --draft-vocab-size to build vocab mappings at training time. Falls back to '<data- path>/token_freq.pt'. If neither exists nor --draft- vocab-size is set, vocab mapping is skipped and the full verifier vocab is used. --draft-vocab-size DRAFT_VOCAB_SIZE Vocabulary size for the draft model. Must be provided together with a token frequency file to generate vocab mappings; otherwise a no-op. --d2t-path D2T_PATH Draft-to-target vocab mapping file (.npy). --t2d-path T2D_PATH Target-to-draft vocab mapping file (.npy). --mask-token-id MASK_TOKEN_ID Token id used to mask positions during training. Resolved from the verifier tokenizer when unset. --norm-before-residual, --no-norm-before-residual Toggle normalization before residual connections (default: True). --embed-requires-grad, --no-embed-requires-grad Whether to train embedding layer weights (default: False). --norm-before-fc, --no-norm-before-fc Apply a single RMSNorm to the concatenated auxiliary hidden states before the FC projection (gpt-oss style). See --fc-norm for the per-layer alternative from the Eagle 3.1 paper. (default: True for eagle3, False otherwise). --fc-norm Apply per-layer RMSNorm to each auxiliary hidden state before concatenation and FC projection (Eagle 3.1 paper approach). --norm-output, --no-norm-output Feed post-norm hidden states back across TTT steps to stabilize magnitude drift across speculation depths (default: True for eagle3, False otherwise). --sliding-window SLIDING_WINDOW Sliding window size for sliding window attention layers (default: 2048). All draft layers use sliding window by default (except mtp). --full-attention-indices FULL_ATTENTION_INDICES [FULL_ATTENTION_INDICES ...] (Optional) space-separated draft layer indices that should use full attention instead of sliding window. All draft layers use sliding window by default (except mtp). (e.g. '--full-attention-indices 0 2' makes layers 0 and 2 use full attention; the rest use sliding window). --sliding-window-non-causal Use non-causal (bidirectional) masking within draft blocks for sliding window attention layers. Full attention layers are always bidirectional. Note: vLLM currently doesn't support these models. --draft-attn-impl {simple_flex_attention,sdpa,eager} Attention implementation for draft layers. Use 'sdpa' or 'eager' for hardware that doesn't support flex attention. Not supported for MTP. data: --data-path DATA_PATH Root data directory with the preprocessed dataset, vocab mappings (d2t.npy, t2d.npy), token frequencies (token_freq.pt), and hidden states. --hidden-states-backend {file} Hidden states transfer backend. Each backend may add its own CLI arguments. Default: 'file'. --hidden-states-path HIDDEN_STATES_PATH Path where cached hidden states are stored (default: <data-path>/hidden_states). Contributed by the 'file' backend. --total-seq-len TOTAL_SEQ_LEN Maximum training sequence length, in tokens. --train-data-ratio TRAIN_DATA_RATIO Fraction of the dataset used for training; the remainder is held out for validation. --noise-std NOISE_STD Standard deviation for noise augmentation. --legacy-data DEPRECATED. Use the old data format which stores hidden states alongside token_ids and assistant_masks in data_i.pt files. Will be removed soon. --hidden-states-dtype HIDDEN_STATES_DTYPE Data type for dataloader hidden states and autocast compute. Model master weights are always kept in fp32. Options: float32, bfloat16 (recommended). float16 is not supported (requires gradient scaling). --num-workers NUM_WORKERS Number of dataloader workers. --prefetch-factor PREFETCH_FACTOR Dataloader prefetch factor. --max-anchors MAX_ANCHORS Maximum anchor positions for DFlash, DSpark, and P-EAGLE training (default: 3072). generation: --vllm-endpoint VLLM_ENDPOINT vLLM endpoint used to generate hidden states on demand. Only needed if --on-missing=generate and samples are missing. The vLLM instance must cache hidden states to a location reachable from the training instance. --on-missing {generate,skip,warn,raise} Dataloader behaviour when a sample has no cached hidden states. Default 'generate' generates them on demand via the vLLM endpoint; the others skip the sample, skip with a warning, or raise. --on-generate {cache,delete} Behaviour after generating a hidden state (only if --on-missing=generate). 'delete' discards it once loaded; 'cache' stores it in the hidden states path, enabling hybrid online/offline training. --request-timeout REQUEST_TIMEOUT Timeout in seconds for each individual vLLM request. Only applies if --on-missing=generate. --max-retries MAX_RETRIES Maximum retry attempts per vLLM request on failure. Only applies if --on-missing=generate. loss: --loss-fn LOSS_FN Loss function specification. A name (kl_div, rkl, jsd, ce, tv, nla, lk_hybrid) or a JSON dict for a weighted combination, e.g. '{"ce": 0.1, "tv": 0.9}'. --ttt-steps TTT_STEPS Number of test-time-training (TTT) steps the draft is unrolled over during training. --ttt-step-loss-decay TTT_STEP_LOSS_DECAY Multiplicative loss weight decay applied per TTT step (1.0 = weight every step equally). optimizer: --optimizer {adamw,muon} Optimizer to use. 'muon' applies Muon to 2D weight matrices and AdamW to the remaining params (norms, biases, embeddings, lm_head). --lr LR Learning rate (AdamW / base group). --weight-decay WEIGHT_DECAY Weight decay for the AdamW optimizer (and the AdamW group in muon mode). --muon-lr MUON_LR LR for the Muon (2D weights) group. Only used with --optimizer muon. Defaults to 10*lr (and --lr defaults to 1e-4). --muon-momentum MUON_MOMENTUM Momentum for the Muon group. --muon-weight-decay MUON_WEIGHT_DECAY Weight decay for the Muon group. --muon-ns-steps MUON_NS_STEPS Newton-Schulz iteration steps for Muon. --muon-adjust-lr-fn {original,match_rms_adamw} Muon LR adjustment. 'match_rms_adamw' matches AdamW's update RMS. scheduler: --scheduler-type {linear,cosine,none} LR scheduler type. --scheduler-warmup-steps SCHEDULER_WARMUP_STEPS Warmup steps (default: scheduler-dependent). --scheduler-warmup-ratio SCHEDULER_WARMUP_RATIO Warmup as a fraction of total scheduler steps, in [0, 1]. Ignored (with a warning) when --scheduler-warmup- steps is also set. --scheduler-total-steps SCHEDULER_TOTAL_STEPS Total scheduler steps (default: inferred). --scheduler-num-cosine-cycles SCHEDULER_NUM_COSINE_CYCLES Number of cosine cycles for the cosine scheduler. trainer: --epochs EPOCHS Number of training epochs. --checkpoint-freq CHECKPOINT_FREQ Save a checkpoint every N epochs. Values < 1 enable sub-epoch checkpointing (e.g. 0.5 = every half epoch). --save-best Also point a checkpoint at the lowest validation loss. --no-resume-from-checkpoint Do not resume training from an existing checkpoint. --log-freq LOG_FREQ Log training metrics every N steps (default: 1). --save-path SAVE_PATH Directory to write checkpoints and the resolved run.yaml. --fsdp-shard Shard model parameters across GPUs with FSDP. By default parameters are fully replicated (DDP-like). Enable when the model does not fit in a single GPU's memory. logging: --logger LOGGER One of 'trackio', 'wandb', 'tensorboard', 'mlflow' or a comma separated list. --log-dir LOG_DIR Directory for training logs. --run-name RUN_NAME Name for this run (used by metric loggers). Auto- generated when unset. dflash: --block-size BLOCK_SIZE Block size for DFlash model (default: 8). --sample-from-anchor, --no-sample-from-anchor Sample from the anchor position (all positions predict). Default: False for dflash, True for dspark. --dflash-decay-gamma DFLASH_DECAY_GAMMA Decay gamma for DFlash/DSpark loss weighting. --per-position-loss-weight {fixed-exp-decay,dpace} Per-position loss weight option for D-PACE support (default: fixed-exp-decay). --dpace-alpha DPACE_ALPHA Smoothing constant for the D-PACE loss (default: 0.5). Must be in (0, 1] when --per-position-loss- weight=dpace. dspark: --markov-rank MARKOV_RANK DSpark: low-rank dim of the Markov logit-bias head (0 disables it). --markov-head-type {vanilla,gated,rnn} DSpark: sequential head variant. --enable-confidence-head, --no-enable-confidence-head DSpark: attach the per-position acceptance confidence head. --confidence-head-with-markov, --no-confidence-head-with-markov DSpark: feed the Markov previous-token embedding into the confidence head alongside the backbone hidden state. --confidence-head-alpha CONFIDENCE_HEAD_ALPHA DSpark: weight of the confidence-head BCE term. peagle: --num-depths NUM_DEPTHS Number of parallel prediction depths for P-EAGLE (default: 8). --down-sample-ratio DOWN_SAMPLE_RATIO Geometric decay ratio for COD sampling in P-EAGLE. --down-sample-ratio-min DOWN_SAMPLE_RATIO_MIN Minimum retention ratio for COD sampling in P-EAGLE. mtp: --num-speculative-steps NUM_SPECULATIVE_STEPS Number of MTP prediction steps (default: 3). --step-weight-beta STEP_WEIGHT_BETA Exponential decay factor for MTP step weights; higher weights earlier prediction steps more. Only used with MTP. ``` </details> - Every existing flag, env, and default is preserved. `flatten()` derives the flat `vars(args)`-shaped dict the model layer consumes directly from the schema (consumers bind by name), so a new tunable is a one-field edit and there is no hand-maintained key list to keep in sync. - The example recipes under `examples/train/*.sh` are replayed through `resolve()` as the central back-compat guard — a renamed/dropped/retyped flag breaks a real recipe. - `hs_connectors` stays speculators-agnostic (argparse-based, usable by vLLM standalone). Backend train-args are mirrored as schema fields; `test_backend_reconciliation.py` fails loudly if a backend registers a train-arg with no matching schema field. - **Reproducibility anchor.** `run.yaml` intentionally emits only non-default values, so it is *not* pinned against future changes to the schema defaults. `train_command.txt` (which records the exact argv **plus the resolving git SHA**, world size, and library versions) is the durable anchor: a reload always names the code version it was produced against. Closes #833 1. **`feat(train/config): typed schema as the single source of truth`** — `schema.py`. The typed groups and the `flatten()`/`from_flat()` seam. `flatten()` derives the flat surface straight from the schema (declaration order), so adding a tunable is a one-field edit. Value validators (e.g. the `hidden_states_dtype` `float16` rejection) run at construction, so bad values fail as clean config errors. 2. **`feat(train/config): generate the CLI from the schema; layered resolution`** — `resolution.py`. Argparse generation from annotations (choices/required/bool markers are declarative via `json_schema_extra`), the `flag > yaml > default` walk, provenance, the post-merge required-flag check, draft-init conflict check, and the backend reconciliation guard. 3. **`feat(train/config): reproducibility artifacts`** — `artifacts.py`. YAML dump + `run.yaml`/`train_command.txt` writers; only non-default values are emitted, in the pinned order, so the artifact round-trips byte-stably (and `dump_yaml()` also handles a `from_flat` config, which carries no provenance). 4. **`refactor(train): cut train.py over to TrainConfig; drop legacy argparse`** — the `scripts/train.py` diff (mostly deletions, −660 net) plus removal of the now-orphaned `argparse_utils` helper. - `pytest tests/unit/train/config/` → **65 passed** (schema / resolution / artifacts / backend-reconciliation seams). - Full `tests/unit/train/` → **221 passed**; the only 2 failures (`test_distributed_resume[ddp]/[fsdp]`) are pre-existing and unrelated — they fail identically on the pre-refactor commit with the same `initial_lr` scheduler-resume error. - Back-compat guard: all six `examples/train/*.sh` recipes replay through `resolve()` cleanly (`-k recipe` → 6 passed). - Round-trip verified: `--dump-config` → reload via `--config` → re-dump is **byte-identical** (regression test: `test_dump_yaml_is_byte_stable_across_reload`). - `ruff check` / `ruff format --check` → clean · `mypy src/speculators/train/config/` → no issues · `scripts/train.py --help` renders the generated grouped help. - Each of the four commits was checked out individually and imports + passes its own tests. - The `argparse.Namespace(**cfg.flatten())` adapter in `main()` is a deliberate phase-1 seam. Phase 2 passes the typed config into the model layer directly and retires the adapter; it is deferred to a separate PR to keep this diff focused on the CLI/resolution boundary. - **Uniform bool override (deferred, intentional).** Only `embed_requires_grad` among the False-default bools is tagged `BooleanOptionalAction`, so a `run.yaml` that sets another False-default bool `true` (e.g. `save_best`, `fsdp_shard`) has no CLI off-switch. Blanket-tagging every False-default bool would make CLI-over-YAML override uniform, but it is *not* free: `no_resume_from_checkpoint` would generate a confusing double-negative `--no-no-resume-from-checkpoint`. Left as a deliberate follow-up rather than shipping that UX wrinkle. - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [x] The test plan/results, such as providing test command and pasting the results. - [x] (Optional) The necessary documentation update. - [x] I (a human) have written or reviewed the code in this pr to the best of my ability. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Rahul-Tuli <rtuli@redhat.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
<!-- markdownlint-disable --> ## Purpose Update response regeneration docs to reflect multi-turn support (#693), retry/resume changes (#716), and new `run_all.sh` flags (#499). ## Description Descriptions, argument tables, and examples in both `response_regeneration.md` files still described single-turn-only behavior. Updates both the tutorial and CLI reference with multi-turn language, a multi-turn JSON example, missing arguments (`--subset`, `--max-retries`, `--max-model-len`, `--reasoning-parser`), and GSM8K in the supported datasets table. ## Tests - `mdformat --check`, `ruff check`, `ruff format --check` pass - All claims verified against `scripts/response_regeneration/script.py` and `run_all.sh` I have filled in: - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [x] The test plan/results, such as providing test command and pasting the results. - [x] (Optional) The necessary documentation update. - [x] I (a human) have written or reviewed the code in this pr to the best of my ability. Signed-off-by: Sawyer Bowerman <sbowerma@redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
## Summary The config-file-first CLI refactor (#841) moved every `scripts/train.py` flag into the typed schema but did not carry over **`--max-steps`**. `TrainerConfig.max_steps` and the training loop's early-stop check (`self.global_step >= self.config.max_steps`) still exist, but the flag was unreachable from the CLI and left unset in `main()`, so `max_steps` was effectively dead code — `--max-steps N` now errors with `unrecognized arguments: --max-steps`. This restores it the schema-native way: - Add a single `max_steps: int | None` `Field` (with `ge=1`) to `TrainerArgs` in `src/speculators/train/config/schema.py`, so the flag, type, `--help`, and validation are all derived from it automatically. - Wire it into the `TrainerConfig` built in `main()` (`max_steps=args.max_steps`). No behavior change to the training loop itself — it already honored `TrainerConfig.max_steps`; this just makes it reachable again. ## Test plan - `python -m pytest tests/unit/train/test_cli_args.py tests/unit/train/config` — 88 passing (adds default/explicit/reject-non-positive cases for `--max-steps`). - `make quality` (ruff + format + mypy) clean. - Manual: `--max-steps 15` → `cfg.trainer.max_steps == 15`; omitted → `None`; `--max-steps 0` → clean config error; flag appears in `--help`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: shanjiaz <zsjwpianpian@gmail.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
#839) ## The issue: The two scripts disagree about what the flag's default is, and the disagreement is correct — each is right for its own side: Script | Default layer ids | Count | Because -- | -- | -- | -- launch_vllm.py | [2, n//2, n−3, n] | 4 | Captures the three auxiliary layers and the final layer n, since both are needed downstream. train.py | [2, n//2, n−3] | 3 | Names only the auxiliary layers. The final layer arrives as a separate input, not as a member of this list. ## This PR: What Docs and help text only; no behavior change. ## Detailed explaination: https://claude.ai/code/artifact/32b15cef-e893-4816-9612-9a072bdb3cb3 ## Tests No behavior change, so no new tests. `ruff check`, `ruff format --check`, and `mdformat --check` pass. `python scripts/train.py --help` renders: ``` --target-layer-ids TARGET_LAYER_IDS [TARGET_LAYER_IDS ...] (Optional) A (space separated) list of integer layer ids for the auxiliary hidden states. Defaults to [2, num_hidden_layers // 2, num_hidden_layers - 3]. Note: if custom values were used to launch vllm, pass the same ids here, excluding the final layer launch_vllm.py appends. ``` ## Checklist I have filled in: - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [x] The test plan/results, such as providing test command and pasting the results. - [x] (Optional) The necessary documentation update. - [x] I (a human) have written or reviewed the code in this pr to the best of my ability. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
…kew (#867) ## Summary - **Sampler rotation**: the LPT bin-packer in `MultipackDistributedBatchSamplerV2` tie-breaks on slot id when bins are equally full, so slot 0 always grabs the largest sample and the highest slot absorbs leftover small ones. Over a full epoch this creates a **2.4× sample-count skew** across ranks. The fix rotates the slot→rank mapping per batch (`target_slot = (rank - rotation) % num_replicas`), cycling which rank plays each role. Sample-count ratio drops from 2.40× to 1.00×; token balance stays at 1.002×; shards remain disjoint; batch counts are unchanged. - **Validation sync barriers**: after PR #813 removed per-batch validation all-reduce, cross-rank speed differences in the validation loop accumulate unboundedly. With the sample-count skew above, the fastest rank finishes validation minutes before the slowest, and the final metrics `all_reduce` trips the NCCL watchdog (10 min default) → **validation hang**. A new `_maybe_val_sync` method issues a `dist.barrier()` every 50 validation batches, bounding drift well below the watchdog window. **Root cause**: LPT slot-id tie-breaking bias → 2.4× val sample-count skew → no per-batch sync (post-#813) → drift exceeds NCCL watchdog → hang at end-of-epoch validation. ## Test plan - [x] 19 unit tests for sampler balance (`test_batch_sampler_balance.py`): sample balance, token balance, disjoint partitions, equal batch counts, token budget, rotation mechanism — 5 of these fail on the unfixed sampler - [x] 8 unit tests for validation sync (`test_val_sync.py`): barrier placement at interval boundaries only, never at batch 0, disabled when single-process or interval=0, deterministic count - [x] Fast-val reproduction: built a dataset with the exact val split (10k rows, 1864 batches/rank) but tiny train split — reproduced the hang on unfixed code, confirmed fix eliminates it - [x] Full 100k-sample ablation (accum=1 and accum=4, 3-GPU DDP, ~8h each): both complete validation without hanging 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
## Summary - The `ArrowDataset` split logic computed the train boundary as `int(n * ratio)` and the val boundary as `int(n * (ratio - 1.0))`. These two expressions disagree in floating point at certain ratios — e.g. at `ratio=0.2` with `n=12501`: `int(12501 * 0.2) == 2500` but `int(12501 * -0.8) → int(12501 * 0.2000…3) == 2501` — putting **one row in both splits** (silent train/val leakage). - The fix replaces the `split_ratio: float` parameter with `train_ratio: float` + `split: Literal["train", "val"]`. Both splits derive their boundary from the single expression `int(len(data) * train_ratio)`, guaranteeing exact complementarity. Input validation rejects out-of-range ratios and val splits at `train_ratio=1.0`. ## Test plan - [x] 34 unit tests (`test_split_boundaries.py`): exact complementarity across 10 ratios (including 0.2 which was broken), no row in both splits, union covers full dataset, default takes whole dataset, rejects out-of-range, rejects val at ratio 1.0 - [x] Verified the 0.2 ratio case produces off-by-one on unfixed code, exact complement on fixed code 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
split_batch_for_sp Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>
Roderick-Wu
force-pushed
the
feat/ulysses-sp
branch
from
July 27, 2026 21:15
982377b to
d346e9d
Compare
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
9 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.
Purpose
Add flag for enabling ulysses sequence parallelism. Reduces memory usage in long contexts.
Set using
--sp-size <number of ranks to shard sequence across>flag.Defaults to using ddp to shard model (replicates, no sharding here), also works with fsdp (
--fsdp-shard).Tests
Tested supported speculator training e2e (eagle3, peagle, dflash, dspark). Verified that memory usage is reduced.
Checklist
I have filled in: