Repository navigation
feat: add Qwen Flash Next training and compact 256K recipes - #1731
dingzhiqiang wants to merge 49 commits into
Conversation
Complete ten fresh 256K steps with identical aggregate reward of 1198/1280 on internal and open-source branches. First-step absolute logp differences are 0.038508 and 0.038055 respectively, at the 1e-2 scale. Record the patched SGLang runtime and remaining numerical limitations.
dingzhiqiang
changed the title
sitabulaixizawaluduo
left a comment
sitabulaixizawaluduo
left a comment
There was a problem hiding this comment.
The added tests mix model-specific contracts with the generic CI image. The unit image uses SGLang 0.5.10 / Transformers 5.3, while this recipe targets SGLang 0.5.19.dev125+g119b5ffe4 and some Qwen4Exp paths require Transformers 5.16.1. As a result, tests either fail in the wrong backend lane or skip the actual model path. Please remove unrelated UTs and avoid treating results from an incompatible runtime image as validation for this feature.
| trackio_config = self.config.trackio | ||
| if trackio_config.mode != "disabled": | ||
| global trackio | ||
| if trackio is None: |
There was a problem hiding this comment.
How is this Trackio import refactor related to the Qwen Flash Next feature? I could not find a call-path dependency from the new recipe or AWEX integration. Bundling an unrelated StatsLogger behavior change increases the regression surface; please move it to a separate PR.
| assert exc_info.value.result.raw == {"error": "grader failed", "arca": {}} | ||
|
|
||
|
|
||
| def test_swe_identity_reward_preserves_partial_scores_and_rejects_invalid_values(): |
There was a problem hiding this comment.
This test covers identity_reward behavior that is unchanged between base and head; reward_transforms.py is not modified by this PR. It therefore does not validate this feature and inflates the reported test coverage. Please remove this unrelated UT.
| name="main", | ||
| split=split, | ||
| cache_dir=_rank_local_cache_dir(cache_dir), | ||
| keep_in_memory=True, |
There was a problem hiding this comment.
This changes the default for every GSM8K caller by forcing the dataset, map, and filter results into memory, even when cache_dir is not configured. The new Qwen RL recipe does not set cache_dir, and the added test only exercises the SFT loader with an explicit cache path, so the global default behavior and RL path are untested. Why is this dataset-wide change part of the model integration?
| drop_incomplete_group=config.gconfig.drop_incomplete_group, | ||
| ) | ||
| if is_v1_rollout: | ||
| prepare_kwargs["reward_normalization_use_std"] = ( |
There was a problem hiding this comment.
This field is read and forwarded on the v1 path even when reward_normalization is false, which changes the default call contract and is the reason existing controller/trainer tests fail. The repository already expresses mean-only group normalization through actor.reward_norm with mean_level=group and std_level=null. Why is another option propagated through every engine/controller needed instead of reusing the existing mechanism?
| if spec is None or spec.origin is None: | ||
| raise RuntimeError("SGLang is not installed; pass --target") | ||
| target = Path(spec.origin).parent / TARGET_RELATIVE | ||
| target.write_bytes(patched_source(target.read_bytes())) |
There was a problem hiding this comment.
The only added test for this patch is the unknown-source rejection case. There is no positive test against the pinned SGLang source, no assertion of PATCHED_SHA256, and no file-write smoke test, so the test suite cannot show that the packaged patch actually applies to the runtime used by the recipe.
| self._model_device = device | ||
| super().__init__(*args, model=model, **kwargs) | ||
|
|
||
| def _init_reader_in_colocate_mode(self): |
There was a problem hiding this comment.
This overrides the AWEX colocate transport for every existing model, not only Qwen4Exp, and defaults expert packing on. The pinned AWEX 0.8.1 already has an AWEX_CHUNK_MB bounded-memory path, while the added repository tests mock P2P and cannot regression-test real NCCL for existing models. Why does a Qwen recipe PR need to replace the generic transport instead of scoping or reusing the upstream path?
|
|
||
|
|
||
| def _expert_pack_limits() -> tuple[int, int]: | ||
| max_ops = _read_positive_env_int("AWEX_EXPERT_PACK_OPS", _DEFAULT_EXPERT_PACK_OPS) |
There was a problem hiding this comment.
AWEX_EXPERT_PACK_OPS and AWEX_EXPERT_PACK_MB change wire batching but are read independently in each process, with no cross-rank agreement check. Different sender/receiver values produce a different number of P2P messages and can hang or corrupt the transfer. These protocol parameters should not be per-process environment overrides; use validated constants or a single negotiated configuration.
| assert_supported_sglang_version() | ||
| from sglang.srt.managers.scheduler import Scheduler | ||
|
|
||
| if os.environ.get("QWEN_AWEX_FROZEN_CONTRACT"): |
There was a problem hiding this comment.
The presence of QWEN_AWEX_FROZEN_CONTRACT is used as both the manifest path and a feature flag. This dedicated recipe already fixes the Qwen AWEX behavior, so activation should be derived from the model/recipe rather than an environment variable; only the model-specific manifest path should remain an explicit input.
| _metadata_gc_lock = threading.Lock() | ||
|
|
||
|
|
||
| def serialize_metadata_gc(func: Callable[_P, _T]) -> Callable[_P, _T]: |
There was a problem hiding this comment.
This process-wide guard changes all AWEX/SGLang schedulers and is not gated to Qwen or the new recipe. It appears to be a separate generic SGLang/AWEX race fix; bundling it here expands the feature scope, and the tests only use fake schedulers.
| _report(violations, "static") | ||
|
|
||
|
|
||
| def resolve_scheduler_parallel_attr(scheduler: Any, name: str) -> int | None: |
There was a problem hiding this comment.
Do we need a separate multi-version fork-contract layer when the plugin already rejects all but explicitly supported SGLang versions and this recipe pins one exact dev revision? The contract has also drifted from the real consumer: it accepts tp_rank from tp_worker/model_runner parallel state, but _patch_execute_task_in_model_worker uses a different lookup and can still default to rank 0. Please consolidate the validation into the actual consumers or remove this file rather than maintaining a parallel contract.
## Description Keep transfer planning as tensor views and allocate contiguous send/receive buffers only for the active batch, including strided column slices. Copy receives back before releasing that batch. Pack compatible expert operations within configured bounds, check packing limits across inference ranks before P2P starts, and provide CUDA IPC staging allocation that restores the allocator settings used for training. The packing byte limit is not a hard cap for a single oversized tensor. Add a reader factory so integrations can explicitly choose the bounded transport; the AWEX default remains unchanged. All participating readers must choose the same transport. This extracts the shared transport implementation from areal-project/AReaL#1731. Own Qwen4Exp weight converters, GDN/gated-QKV layouts and frozen-parameter declarations in AWEX. Registration is explicit and accepts caller-supplied binders; checkpoint evidence and live engine ownership stay in the integration. AReaL consumes these APIs in areal-project/AReaL#1751 without a second converter implementation. ## Validation - 41 targeted CPU tests passed, including real two-rank Gloo agreement/disagreement, strided buffer allocation/release and independent GDN/QKV coordinate checks across TP sizes. - AReaL integration regression suite: 557 passed, 9 optional-runtime skips; separate Transformers 5.16.1 vision run: 31 passed. - Ruff check and format passed for changed files. - GPU NCCL transfers, repeated weight equality, and cross-image CUDA IPC remain unverified and require validation before merge.
Description
Add Qwen Flash Next training integration, AWEX weight exchange, and compact SWE/GSM8K/SFT recipes. Qwen-specific PLE checkpoint validation, fixed-weight export completion, and export configuration live in mcore-bridge; AReaL retains the engine adapter, training parameter policy, and distributed checkpoint coordination. Remove the global Megatron configuration-validation bypass.
Depends on modelscope/mcore-bridge#202, pinned to
38c675ed356c3ab316da7ba9a88857952591bbc6in both dependency configurations and lockfiles. No Megatron-LM implementation is vendored.The recipes include two source-hash-guarded SGLang QSA patches and apply them before rollout workers start. They require the specified base inference image with a fresh writable container layer and an externally supplied model-specific frozen-weight manifest. Experiment diagnostics and fixed task fixtures are excluded.
Final review fixes synchronize AWEX slice preparation before transfer-stream execution, including ranks without local self-copy, and allocate contiguous receive scratch without copying stale weights. The pinned bridge validates complete PLE dequantized exports in BF16, FP16, or FP32 while keeping non-PLE dtype checks strict.
SWE alignment: internal vs. open-source
Matched fresh 10-step SWE training experiments on 64 GPUs, using the same model, 16-task subset, batch size, seed, learning rate, and parallelism:
Aggregate training reward is identical; first-step
ppo_actor/update/logp_abs_diff/avgdiffers by approximately 1.18%. Both runs completed without OOM. This is a training acceptance subset, not an independent SWE-bench evaluation. These alignment runs precede the final checkpoint split and do not establish identical trajectories or exact numerical equivalence.Checkpoint-split runtime validation
AReaL
b0c3ff41fwith bridge4964a6fcompleted two fresh 64-GPU SWE training steps using TP8/PP8/EP8/CP1 training and TP4 inference.Both final-run AWEX weight updates completed on all 64 actor ranks. Step 2 had nonzero gradient norm 0.28691. Final HF checkpoint export completed: 1,627 indexed tensors across 121 safetensors files, with index entries and file sizes checked; all 128 PLE table shards were contiguous in numbering and BF16, with obsolete FP8 scales removed. Controller exited successfully and all workers were released.
256K PLE indexing regression
A full 32-group × 8-sample memory probe on 64 GPUs generated all 256 sequences at 262,142 tokens without observed inference OOM, then failed during PPO with a CUDA illegal memory access. A single-GPU reproduction isolated signed int32 overflow in the bridge PLE kernels: 262,144 tokens × 10,240 features exceeds 2^31 elements. The pinned bridge promotes indices to int64 before multiplication.
The old gate kernel reproduced the failure; the fixed kernel passed. Eight GPU regression tests passed: six forward/backward kernels at the 256K address boundary and two small nonzero forward/all-input-gradient comparisons against PyTorch. Full pre-commit checks passed in both repositories. Fixed-length 256K and normal variable-length SWE runs are being repeated with the fix; complete PPO training without OOM is not yet established.
Validation
5abc9cc03; 196 targeted tests passed at that revision. Read-only Megatron review confirmed the migrated function bodies and signatures matched their original implementations.Runtime limitations
Validated inference runtime: SGLang
0.5.19.dev125+g119b5ffe4with the included stable top-k and compress-gather fixes. Already patched or different sources are rejected. Patch-free upstreamd72e595remains unvalidated. AWEX 0.8.1 and Megatron-Core 0.17.0 are separate dependencies. Historical post-update probe discrepancies remain unresolved.The checkpoint-split revisions listed above completed multi-node CP1 SWE validation. CP>1 runtime validation and a full checkpoint recovery run remain unverified. Full training validation of the latest PLE indexing fix is pending; the 256K configuration ceiling does not establish that every SWE trajectory reaches 256K.
Created and updated by a coding agent.
256K QSA indexing follow-up
Further validation reproduced a separate signed int32 overflow in the QSA selection bitmap (262,144 × 65,536 bytes). Bridge
38c675epromotes indices before multiplication in forward, dQ, and dK/dV, including indices loaded from the sparse tile lists. The old kernel failed with illegal memory access; the fixed kernel passed forward and backward on the same 256K shape.All 11 PLE/QSA GPU regression tests passed, including nonzero dense-reference forward/gradient comparisons and actual 256K address-boundary tests. Full pre-commit checks passed in both repositories. AReaL
097e146c9pins this bridge. The PLE-only experimental reruns were canceled after this second bug was reproduced; fixed-length 256K (32 × 8, 64 GPUs) and normal variable-length SWE (same batch and hardware, two steps) are being resubmitted with both fixes. End-to-end training without OOM remains pending.