Repository navigation
fix(module-inject): preserve HF precision boundaries in segKI kernels - #9057
Draft
nathon-lee wants to merge 43 commits into
Draft
nathon-lee wants to merge 43 commits into
nathon-lee wants to merge 43 commits into
Conversation
…ake, FLA in-place recurrent state Squashed from 20 commits of the A100 optimization session (net effect): - Ampere GLU routes to a concat-cached cuBLAS GEMM plus a native bf16 silu_mul kernel (the WMMA fused kernel stays for Ada and small M); +4% at bs<=32 on A100 TP2 - custom all-reduce gains a bounded publish-and-wait handshake so fabrics where the flag protocol does not propagate (PCIe P2P hosts) fall back to NCCL instead of hanging - GDN recurrent state streams from gmem in our in-place kernel (the 64KB smem staging only capped occupancy) — kept for the TP=1 path - FLA's Triton recurrent kernel is called directly with the cache slot as both h0 and ht: the wrapper's fresh allocation and the per-layer copy-back disappear, and decode moves off the pure-torch fp32 reference that transformers 5.x falls back to (+9-15% TP2; bypass is bitwise-identical to FLA's public API; FLA_BYPASS=0 restores the wrapper) - lab notes: INVESTIGATIONS.md items 1-3 and the FLA-bypass unraveling Baseline 963/2756/4255 -> 1042-1050/3108-3113/4859-4889 tok/s at bs 8/32/64 (2xA100 SXM4 TP2), 71/64/64% of vLLM 0.32. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…atch files out - bf16 WMMA kernel bodies guard on __CUDA_ARCH__>=800 (precompile builds every arch; sm75 compiles an empty shell that is never launched) - torch.distributed literals -> deepspeed.comm (segment_ki, layers); custom_ar.py marks its deliberately CUDA-specific calls #ignore-cuda; INVESTIGATIONS.md rewords the literal check-torchdist greps - padded-prompt guard reads module.config defensively: test doubles without one crashed the shared-prefill suite - session working files (experiments/, review notes, diff dumps) drop out of the branch — they tripped whitespace hooks and added review noise Verified on 2x GPU (AutoDL): module_inject + rollout suites 271 passed, 2 skipped, 0 failed. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
… markers, drop scratch notes Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…ight caches In-place weight writes (arena syncs, reshard materializations, training steps) keep weight addresses, so the injected forwards' derived caches (concat groups, sliced conv weights, fp32 gate constants, the backend's pointer-keyed GLU concat) would serve stale copies mixed with fresh direct reads. segki_refresh(model) — also exposed as HybridEngineRollout.refresh_weights() — drops them all; a fresh injection rebuilds the same state. Contract tests: after an in-place update + refresh, generation is BITWISE equal to a fresh injection over the same weights, including when parameter storage is bound to external arena tensors updated in place. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Two real AdamW steps on a ZeRO-3 engine (world=2); gathered weights written in place into the TP=2 rollout's sharded storage (split dim inferred from shape mismatch — the storage-map reshard contract); segki_refresh; generate must be bitwise equal to a fresh injection. Validated on 2x4080 SUPER with Qwen3.5-0.8B-Base: PASS both ranks, 0 token diffs. Opt-in via DS_RUN_ZERO3_E2E=1. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…dshake design Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…torch.cuda Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…nse headers Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…op the local custom_ar The one-shot P2P all-reduce now lives in its own PR with a cleaner shape: builder-based load behind the accelerator contract, lazy initialization inside the inference_all_reduce hook (no explicit init call at rollout construction), per-op readiness with a collective failure agreement, and a simplified RowParallel routing that delegates all eligibility checks to the hook. This PR keeps only the consumer side; the implementation carries over verbatim from deepspeedai#8997. - remove csrc/inference/custom_ar.cu + deepspeed/comm/custom_ar.py - torch.py/layers.py route through custom_collective (8997's text) - drop the now-dead RowParallel._rs_ag - e2e script needs no explicit init anymore (lazy in the hook) Signed-off-by: Guokai Ma <guokai.ma@intel.com>
silu_mul (the Ampere GLU epilogue) and repeat_qk (the scan head-repeat) gained kernel_reference entries and direct reference-matched tests — they had only indirect coverage through grouped_linear's matrices and the FLA end-to-end. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The torch.library registration (segki_torchlib.py) and the opt-in torch.compile wrapper were a zero-value experiment: the GDN forward's Python-object dependencies broke the compile graph before any fusion, measured -2.3%. segKI's identity is hand-written, per-architecture tuned kernels behind the portable opbuilder interface; glue fusion via compile belongs to DeepCompile, which already uses the same custom-op pattern and can integrate segKI kernels when that work happens. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The check catches real dataloader bugs in training; removing it was a workaround for broken NCCL infrastructure (specific RunPod instances), not a fix for a DeepSpeed defect. On healthy infrastructure the check passes normally. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
… opbuilder interface The internal GEMV/GEMM kernels (single/triple/quad_gemv, single/triple/ quad_gemm, dual_gemv_silu_mul) were exposed via m.def and called directly from segment_ki.py — a historical artifact from before grouped_linear existed as the unified dispatcher. Now: - segment_ki.py: all call sites go through grouped_linear (the flat/ gemv/wmma/composite routing is the dispatcher's job) - fused_glu.cu: 6 internal m.def entries removed; kernels remain as static host functions serving the dispatcher only - m.def count: 17 -> 10 (the portable interface + cache management) Other backends implementing the opbuilder contract see exactly the 10 methods they need; nothing else leaks through the boundary. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…tion The reference-consistency tests were calling the removed internal kernels directly; now all go through the portable interface. One test was asserting on a stale pre-allocated tensor instead of the actual kernel output (caught by the interface cleanup). Signed-off-by: Guokai Ma <guokai.ma@intel.com>
12 functions → 7: the 6 special-path references (dual_gemv_silu_mul, triple_gemv, triple_gemm, quad_gemm, dual_gemm_silu_mul, gdn_input_proj) collapse into a single grouped_linear(x, weights, glu_epilogue) that covers all routing paths. Other backends now read exactly 7 functions to understand the full contract. Test call sites updated accordingly; the segment_ki fallback paths route through the unified reference. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
… interface kernel_reference.py collapses to 7 interface-aligned functions; test and segment_ki fallback call sites updated to the unified grouped_linear entry; _as_concat helper normalizes the CUDA's per-path output format (n=4 flat concatenates, others return separate). Verified: 24/24 passed on GPU. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
q sliced from the fused q|k|v row block (fused[:, :qn]) keeps the fused row stride; the kernel stepped rows by num_q_heads*q_hs (packed layout), so at batch>1 every row after the first read k/v memory as the query. Pass the real row stride into the kernel and fail-closed on a gate whose row stride disagrees. Verified on L20 (2x rows, identical prompts): strided-vs-contiguous maxdiff 0.0000 both rows; b=1 e2e regression token-exact. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The raw-kernel launch skipped FLA's input_guard, so composite-path decode q/k/v (strided views of the fused projection, batch stride = full fused width) were read with packed-stride pointer math — every batch row after the first read k-region memory (o mismatch 0.29; L20, fla 0.6.0). The launch now packs its inputs (skip-if-contiguous), and a one-time parity probe against the wrapper on cloned live tensors gates the bypass per shape, silently rejecting FLA builds whose kernel grid contract differs (0.5.2 reads a 2-D grid). Verified b=2 eager/graph token-exact, b=1 e2e regression match; probe overhead is once per process and shape. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The native (non-segKI) GDN forward fetches conv/recurrent state through the cache's top-level get_conv_state/get_recurrent_state and branches on is_recording_past; only the hybrid slots' attribute forwarding existed, so use_graph_capture=True with use_segki=False (the default) crashed on hybrid models. Delegate the three entries to the bound slots (behavioral query answers False on KV layers); pre-5.19 attribute-style access was already served by the slot forwarding. Covered by a CPU protocol unit test and a GPU graph-without-segKI parity test (token-exact vs generate); found by the decode benchmark baseline, which went from crash to a full batch sweep. (commit hook flake8 skipped: sandbox blocks its multiprocessing semaphores; verified standalone on the same files) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…plementation Pick deepspeedai#8997's final state for the four files that had drifted: kernel link-safety + clang-format (csrc), the CPU-accelerator route exclusion (torch.py), RowParallel's inference_all_reduce-owned backend selection (layers.py), and the restructured op builder. The Python wrapper and the parity test were already identical. (commit-hook flake8 skipped: sandbox blocks its multiprocessing semaphores; each file passes standalone) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…te replication Under AutoTP a custom layer_specs list REPLACES the built-ins; partial configs left the GDN input-side weights accidentally replicated while the TP forward's endpoint arithmetic assumed sharded widths — the a slice came out 2096 rows wide (expected 16) and gdn_gates read past the buffer (L20, Qwen3.5-4B, GVA v/k=2). The TP branch now slices this rank's head range out of the replicated z/b/a blocks alongside qkv and conv1d, and installation fails loudly when any GDN input projection is found sharded (the contract: replicated in_proj_*/conv1d, row-parallel out_proj, standard column/row for attn+MLP). Verified TP=2 b=1 vs TP=1: token-exact over 32 generated tokens; 143.6 tok/s decode (vs vLLM 0.32 TP2 141.4 on the same L20). Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…e_from_logits Master's deepspeedai#8939 adds _token_selector + _repetition_penalty + _apply_repetition_penalty using HF's official warpers and matching generate()'s distribution exactly. It subsumes our hand-rolled _sample_from_logits (18 lines, no repetition penalty). All segKI integration preserved (graph capture, GDN handling, decode attention install); only the sampling mechanism swapped. (commit-hook flake8 skipped: sandbox blocks its multiprocessing semaphores; the file passes standalone) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The merge with deepspeedai#8939 left two copies of get_query_offset on DeepSpeedStaticCache (ours from the pre-merge tree and master's from deepspeedai#8852); keep master's richer per-layer version. The test's DSStaticGDNSlot import became unused when the assertion moved to cache-level methods. (flake8 skipped: sandbox semaphores) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The merge with deepspeedai#8939 adopted _token_selector for the b=1 full-step graph path but two auto-merged regions kept the old hardcoded argmax: the first token after prefill and every step of the b>1 graph+Python loop. With temperature>0, all 200 sampled first tokens were identical (the argmax) while HF generate spread over 26 tokens — the distributions diverged from step 0. Both sites now use select_next_tokens, matching master's deepspeedai#8939 behavior. (flake8 skipped: sandbox semaphores) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…pSpeedStaticLayer get_seq_length() counts the slot the current query is writing (one past the write position) per the upstream deepspeedai#8852 protocol; position and mask math goes through get_query_offset(), which returns the write position itself. Our decode paths are unaffected: they pass cache_position and position_ids explicitly, so this value never feeds RoPE. Also drop the shadowed duplicate container get_query_offset (dead code — deepspeedai#8852's definition wins at class creation). Fixes the cpu-torch-latest failures in test_static_cache_query_offset_precedes_current_write. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The static-layer get_seq_length was restored to write_position + 1 (transformers >= 5.20 protocol, commit 39210bf), but three tests still expected the old no-+1 convention. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The merge with deepspeedai#8980 introduced a local use_graph_capture variable gated by is_graph_capture_supported(), but the routing condition kept self.use_graph_capture (the raw config value). When a mock accelerator reports no graph support, _generate_graph was still selected instead of falling back to eager generate. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…_greedy Adopt master's routing structure exactly (deepspeedai#8980 + deepspeedai#8939): the local use_graph_capture is the config value narrowed by is_graph_capture_supported(). No is_greedy gate — deepspeedai#8939 added sampling to the graph decode path (token selection runs outside the captured graph, ~0.9% decode cost). segKI's only increment beyond master is the needs_pad_fallback condition. The previous commit referenced the local variable without defining it. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
Master's deepspeedai#8939 refactored _continuous_next_tokens to accept a select_next_tokens parameter (applying repetition penalty via _apply_repetition_penalty, then delegating token selection to the injected selector). Our auto-merged version kept the old 5-arg inline implementation. Updated the method, both call sites, and added select_next_tokens definitions in the enclosing continuous methods. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The previous commit added @classmethod without removing the existing @staticmethod; both together make the method uncallable as a bound method. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
Master's tests mock the module-level logger and use SimpleNamespace mocks that lack num_attention_heads. Our merge dropped the logger definition and the TP head-count adjustment assumed the attribute exists. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…iction Same auto-merge pattern: master's deepspeedai#8939 added a sampling parameter to _continuous_prefill_cache and removed the greedy-only restriction in continuous batching. Both call sites updated. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The previous commit added sampling to _continuous_prefill_cache's signature and its call sites inside _continuous_prefill, but forgot to add sampling to _continuous_prefill itself and its callers — yielding F821 undefined name at the inner call sites. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The deepspeedai#8939 merge resolution dropped three master behaviours in hybrid_engine_rollout.py: - _generate_graph lost the flash_attention_2 rejection, so graph capture proceeded into StaticCache and crashed inside transformers instead of raising the documented ValueError. - _generate_continuous routed on self.use_graph_capture directly, losing the accelerator capability gate, so backends without the graph-capture opt-in entered the graph path instead of falling back to eager continuous batching. - generate() logged through a local logging.getLogger while DeepSpeed's logger is the one exposing warning_once. Restore master's exact code for all three and import the DeepSpeed logger. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…oups The CPU shm inference_all_reduce_ kernel is initialized with the WORLD topology and takes no group argument, so it can only produce a correct sum for a world-wide reduction. When a proper TP subgroup called inference_all_reduce (autotp_size=2 inside a 4-rank world), the kernel combined the two sibling subgroups' partial sums into a 4-way sum, corrupting outputs and grad norms. This is the cpu-torch-latest-v1-multirank CI failure: TestTpGradNorm [2-1]/[2-2] failed deterministically (base_norm 1.6128 vs tp_norm 2.0052, bit-identical across runs) on avx512bw runners where the shm kernel engages; avx2-only runners silently fell back to gloo and passed, which made the check look flaky run-to-run. Fix: engage the shm kernel only when the requested group spans the whole world; a subgroup with full world size is necessarily the world itself, so a size comparison suffices. Every other group falls through to the group-aware torch.distributed.all_reduce. No other backend or routing change: the accelerator custom-collective path is already bound to its registered group per process and is untouched. Add a regression test that reduces through two 2-rank subgroups of a 4-rank world and asserts subgroup-scoped results (1/5, not the world-sum 6). With the guard removed the test reproduces the CI failure exactly. Validated on CPU (torch 2.12.1+cpu, python 3.11): - unit/v1/autotp/test_autotp_training.py -k TestTpGradNorm: 4 passed, 2 skipped (previously 2 failed, bit-identical to CI) - unit/comm/test_custom_collective.py: 5 passed, 2 skipped - unit/v1/autotp/test_tp_plan_e2e.py: test_tp_plan_basic_training passed - full test_autotp_training.py: no regressions vs the unpatched head Signed-off-by: Ma, Guokai <guokai.ma@gmail.com>
Signed-off-by: nathon-lee <leejianwoo@gmail.com> fix(module-inject): preserve HF precision boundaries in segKI kernels Signed-off-by: nathon-lee <leejianwoo@gmail.com>
tohtana
added
the
needs-gpu-test
tohtana
removed
the
needs-gpu-test
This branch has not been deployed
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
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.
Partial numerical-correctness follow-up to #8602 at
2d8f4094b3e357f3fe849a789749b193538f7eeb, related to #9017. This does not resolve batch=8 graph parity and is not merge-ready.Changes
--use_fast_mathfrom the fused-glu builder.Sibling checks cover GLU GEMV/WMMA/standalone SiLU, portable references, both GDN gate dtypes and partial/full RoPE. No graph-runtime, ZeRO, optimizer, config or checkpoint changes are included.
Results
Qwen3.5-0.8B-Base, BF16/SDPA, seeds 0/1, batches 1/2/4/8, 128 prompt tokens, 16 greedy output tokens, EOS disabled. Each configuration has one eager check and two graph checks, with 15 actual replays per graph run.
The tested eager batch=8 and graph batch=4 mismatches are corrected. Seed=0, batch=8 graph still fails twice against default HF, unchanged from baseline. HF's own uncompiled static-cache generation also diverges without DeepSpeed injection.
A restored-state lockstep diagnostic found bitwise-identical logits, KV caches and GDN states between real graph replay and same-path eager execution for all 15 steps, both with and without segKI. Probed norm differences from HF were one BF16 step, reproduced by an independent FP32 reduction tree; this does not explain or validate every remaining cross-path difference.
Focused tests: 27 passed, 2 failed (graph batch=8, segKI on/off). The existing rollout/train/rollout test passes separately. Changed-file pre-commit and
git diff --checkpass; failing assertions remain active.Scope and remaining work
Please treat this as a proposed contribution to #8602, not a competing implementation. The 2026-11-11 duplicate check found no open PR mentioning
9017in its body, but the existing segKI PR is still open; any separate PR needs coordination.Please confirm whether acceptance requires exact tokens against HF's default path or same-path graph/eager consistency with numerical tolerances. Keep that discussion in #9017 before changing assertions or scoping a follow-up; no second PR is assumed.
Validated only on single-GPU BF16 greedy inference, CPU composite forward and minimal backward/train-rollout tests. Full ZeRO/offload, distributed parallelism, other precisions/devices and compile paths are unvalidated. Throughput, peak memory and convergence are unmeasured; removing fast math may affect performance.
AI assistance was used. Human line-by-line review and relevant test execution are required before submission; no new PR has been opened.
Environment and commands run
RTX A5000 x1; driver 595.91.07; CUDA toolkit 12.8; Python 3.12.3; PyTorch 2.9.1+cu128; Transformers 5.22.2; Triton 3.5.1; flash-linear-attention 0.6.0. No causal-conv1d; HF uses reference PyTorch convolution. Commands ran from the checkout in an isolated environment with model/JIT caches.
The 24-check matrix used session-local
reproduce.py; the additional diagnostics usedgraph_eager_lockstep.pyandtrace_norm_rounding.py. These scripts are not included in the patch. The repository regression test covers seed=0, batches 1/4/8, segKI on/off and up to two graph attempts; it is not the full matrix above.