ulysses sp draft#853
Conversation
|
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:
✨ 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 |
df5077a to
3b01abb
Compare
Signed-off-by: Roderick Wu <Roderick-Wu@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
3b01abb to
62b0b5e
Compare
Signed-off-by: Roderick-Wu <rowu@redhat.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Roderick Wu <roderickwu2003@gmail.com>
Signed-off-by: Roderick-Wu <rowu@redhat.com>
also block tracing through _AllToAllSP Signed-off-by: Roderick-Wu <rowu@redhat.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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
| maybe_setup_distributed(sp_size=args.sp_size) | ||
|
|
||
| if args.sp_size > 1: | ||
| if args.total_seq_len % args.sp_size != 0: | ||
| raise ValueError( | ||
| f"--total-seq-len ({args.total_seq_len}) must be divisible " | ||
| f"by --sp-size ({args.sp_size})" | ||
| ) | ||
| if args.speculator_type not in ("eagle3", "dflash"): | ||
| raise ValueError( | ||
| f"Sequence parallelism (--sp-size > 1) is currently only " | ||
| f"supported for eagle3 and dflash, " | ||
| f"got --speculator-type={args.speculator_type}" | ||
| ) | ||
| if args.speculator_type == "dflash" and args.max_anchors % args.sp_size != 0: | ||
| raise ValueError( | ||
| f"--max-anchors ({args.max_anchors}) must be divisible " | ||
| f"by --sp-size ({args.sp_size}) for DFlash SP" | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Validate --sp-size against actual distributed launch state.
These checks only look at args.sp_size, never at whether distributed was actually initialized. If a user passes --sp-size 2 without torchrun (no LOCAL_RANK), maybe_setup_distributed returns early and get_sp_size() stays 1 (see distributed.py's early return), so all the divisibility/type checks here still run and pass/fail on unrelated grounds, and training then silently proceeds without sequence parallelism — the user's --sp-size request is silently ignored. Compare with the existing --fsdp-shard guard right below (if args.fsdp_shard and not is_distributed(): raise ValueError(...)); an equivalent guard is missing for sp_size.
🐛 Suggested fix
if args.sp_size > 1:
+ if not is_distributed():
+ raise ValueError(
+ "--sp-size > 1 requires launching with torchrun/distributed "
+ "training; otherwise sequence parallelism has no effect."
+ )
if args.total_seq_len % args.sp_size != 0:Based on path instructions for scripts/**/*.py to "check that scripts handle argument parsing robustly," this validation gap should be closed.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| maybe_setup_distributed(sp_size=args.sp_size) | |
| if args.sp_size > 1: | |
| if args.total_seq_len % args.sp_size != 0: | |
| raise ValueError( | |
| f"--total-seq-len ({args.total_seq_len}) must be divisible " | |
| f"by --sp-size ({args.sp_size})" | |
| ) | |
| if args.speculator_type not in ("eagle3", "dflash"): | |
| raise ValueError( | |
| f"Sequence parallelism (--sp-size > 1) is currently only " | |
| f"supported for eagle3 and dflash, " | |
| f"got --speculator-type={args.speculator_type}" | |
| ) | |
| if args.speculator_type == "dflash" and args.max_anchors % args.sp_size != 0: | |
| raise ValueError( | |
| f"--max-anchors ({args.max_anchors}) must be divisible " | |
| f"by --sp-size ({args.sp_size}) for DFlash SP" | |
| ) | |
| maybe_setup_distributed(sp_size=args.sp_size) | |
| if args.sp_size > 1: | |
| if not is_distributed(): | |
| raise ValueError( | |
| "--sp-size > 1 requires launching with torchrun/distributed " | |
| "training; otherwise sequence parallelism has no effect." | |
| ) | |
| if args.total_seq_len % args.sp_size != 0: | |
| raise ValueError( | |
| f"--total-seq-len ({args.total_seq_len}) must be divisible " | |
| f"by --sp-size ({args.sp_size})" | |
| ) | |
| if args.speculator_type not in ("eagle3", "dflash"): | |
| raise ValueError( | |
| f"Sequence parallelism (--sp-size > 1) is currently only " | |
| f"supported for eagle3 and dflash, " | |
| f"got --speculator-type={args.speculator_type}" | |
| ) | |
| if args.speculator_type == "dflash" and args.max_anchors % args.sp_size != 0: | |
| raise ValueError( | |
| f"--max-anchors ({args.max_anchors}) must be divisible " | |
| f"by --sp-size ({args.sp_size}) for DFlash SP" | |
| ) |
🤖 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 `@scripts/train.py` around lines 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.
Source: Path instructions
| parser.add_argument( | ||
| "--sp-size", | ||
| type=int, | ||
| default=1, | ||
| help=( | ||
| "Sequence parallelism degree (Ulysses). Splits the packed " | ||
| "sequence across this many ranks within each SP group, " | ||
| "reducing per-device activation memory. Requires " | ||
| "world_size %% sp_size == 0 and total_seq_len %% sp_size == 0. " | ||
| "Currently only supported for eagle3." | ||
| ), | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale help text: --sp-size also supports dflash, not just eagle3.
The help string says "Currently only supported for eagle3" but main() (lines 534-539) explicitly allows speculator_type in ("eagle3", "dflash").
📝 Suggested fix
- "Currently only supported for eagle3."
+ "Currently only supported for eagle3 and dflash."📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| parser.add_argument( | |
| "--sp-size", | |
| type=int, | |
| default=1, | |
| help=( | |
| "Sequence parallelism degree (Ulysses). Splits the packed " | |
| "sequence across this many ranks within each SP group, " | |
| "reducing per-device activation memory. Requires " | |
| "world_size %% sp_size == 0 and total_seq_len %% sp_size == 0. " | |
| "Currently only supported for eagle3." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--sp-size", | |
| type=int, | |
| default=1, | |
| help=( | |
| "Sequence parallelism degree (Ulysses). Splits the packed " | |
| "sequence across this many ranks within each SP group, " | |
| "reducing per-device activation memory. Requires " | |
| "world_size %% sp_size == 0 and total_seq_len %% sp_size == 0. " | |
| "Currently only supported for eagle3 and dflash." | |
| ), | |
| ) |
🤖 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 `@scripts/train.py` around lines 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().
| 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) | ||
|
|
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
| # 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] |
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.
| out: dict[str, torch.Tensor] = {} | ||
| full_seq_len: int | None = None | ||
|
|
||
| for key, tensor in batch.items(): | ||
| if key in _SPLIT_KEYS and tensor.dim() >= _MIN_SPLIT_NDIM: | ||
| seq_len = tensor.shape[1] | ||
| if full_seq_len is None: | ||
| full_seq_len = seq_len | ||
| chunks = tensor.chunk(sp_size, dim=1) | ||
| out[key] = chunks[sp_rank].contiguous() | ||
| elif key in _KEEP_FULL_KEYS: | ||
| if full_seq_len is None and tensor.dim() >= _MIN_SPLIT_NDIM: | ||
| full_seq_len = tensor.shape[1] | ||
| out[key] = tensor | ||
| else: | ||
| out[key] = tensor | ||
|
|
||
| if full_seq_len is not None: | ||
| out["full_seq_len"] = torch.tensor(full_seq_len, device=tensor.device) | ||
|
|
||
| return out | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
full_seq_len device inferred from a loop-leaked variable, not the actual split tensor.
tensor at line 116 is whatever the last item in batch.items() happened to bind to in the loop — not necessarily a tensor related to full_seq_len, and not guaranteed to even be a torch.Tensor (batches can carry non-tensor metadata, e.g. lengths, per the trainer's kwargs handling). This is fragile: it works only by coincidence that all batch tensors share a device today, and can raise AttributeError: 'X' object has no attribute 'device' if the last dict entry isn't a tensor.
🐛 Suggested fix
out: dict[str, torch.Tensor] = {}
full_seq_len: int | None = None
+ full_seq_len_device: torch.device | None = None
for key, tensor in batch.items():
if key in _SPLIT_KEYS and tensor.dim() >= _MIN_SPLIT_NDIM:
seq_len = tensor.shape[1]
if full_seq_len is None:
full_seq_len = seq_len
+ full_seq_len_device = tensor.device
chunks = tensor.chunk(sp_size, dim=1)
out[key] = chunks[sp_rank].contiguous()
elif key in _KEEP_FULL_KEYS:
if full_seq_len is None and tensor.dim() >= _MIN_SPLIT_NDIM:
full_seq_len = tensor.shape[1]
+ full_seq_len_device = tensor.device
out[key] = tensor
else:
out[key] = tensor
if full_seq_len is not None:
- out["full_seq_len"] = torch.tensor(full_seq_len, device=tensor.device)
+ out["full_seq_len"] = torch.tensor(full_seq_len, device=full_seq_len_device)As per path instructions for src/speculators/train/**/*.py to verify "device placement is consistent," this should use the device from the tensor that actually determined full_seq_len.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| out: dict[str, torch.Tensor] = {} | |
| full_seq_len: int | None = None | |
| for key, tensor in batch.items(): | |
| if key in _SPLIT_KEYS and tensor.dim() >= _MIN_SPLIT_NDIM: | |
| seq_len = tensor.shape[1] | |
| if full_seq_len is None: | |
| full_seq_len = seq_len | |
| chunks = tensor.chunk(sp_size, dim=1) | |
| out[key] = chunks[sp_rank].contiguous() | |
| elif key in _KEEP_FULL_KEYS: | |
| if full_seq_len is None and tensor.dim() >= _MIN_SPLIT_NDIM: | |
| full_seq_len = tensor.shape[1] | |
| out[key] = tensor | |
| else: | |
| out[key] = tensor | |
| if full_seq_len is not None: | |
| out["full_seq_len"] = torch.tensor(full_seq_len, device=tensor.device) | |
| return out | |
| out: dict[str, torch.Tensor] = {} | |
| full_seq_len: int | None = None | |
| full_seq_len_device: torch.device | None = None | |
| for key, tensor in batch.items(): | |
| if key in _SPLIT_KEYS and tensor.dim() >= _MIN_SPLIT_NDIM: | |
| seq_len = tensor.shape[1] | |
| if full_seq_len is None: | |
| full_seq_len = seq_len | |
| full_seq_len_device = tensor.device | |
| chunks = tensor.chunk(sp_size, dim=1) | |
| out[key] = chunks[sp_rank].contiguous() | |
| elif key in _KEEP_FULL_KEYS: | |
| if full_seq_len is None and tensor.dim() >= _MIN_SPLIT_NDIM: | |
| full_seq_len = tensor.shape[1] | |
| full_seq_len_device = tensor.device | |
| out[key] = tensor | |
| else: | |
| out[key] = tensor | |
| if full_seq_len is not None: | |
| out["full_seq_len"] = torch.tensor(full_seq_len, device=full_seq_len_device) | |
| return out |
🤖 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/sequence_parallel.py` around lines 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.
Source: Path instructions
eagle3 and dflash and does actually reduce memory usage. Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
1 participant
Purpose
Tests
Checklist
I have filled in: