Modernize: VeOmni 8ca09d7 migration + LLaDA2.1/2.2 support + tests - #1
Open
Akicou wants to merge 12 commits into
Open
Modernize: VeOmni 8ca09d7 migration + LLaDA2.1/2.2 support + tests#1Akicou wants to merge 12 commits into
Akicou wants to merge 12 commits into
Conversation
…sync findings Phase 0 (inventory): environment reality (Win11/py3.14, no torch/veomni, RTX 5070 Ti), git state, VeOmni pin (600fe6d, 495 commits behind main), architecture map, and the model-identity branch points Phase 4 depends on. Phase 1 (upstream sync): fork is 1 ahead / 0 behind upstream/main; no upstream branch/tag adds LLaDA2.1/2.2. Records upstream open PR inclusionAI#22 (veomni-npu-migration) as the approved basis for Phase 2.
…migration Adopts the validated upstream migration 'veomni-npu-migration' (PR inclusionAI#22) as the Phase 2 basis (user-approved). Brings: - VeOmni submodule bump 600fe6d (2025-10) -> 8ca09d7 (2026-06), ~9 months. - New nested VeOmni arg schema (VeOmniArguments; train.accelerator/optimizer/ checkpoint/gradient_checkpointing/wandb, data.dataloader, model.ops_implementation). - Entrypoints collapsed into tasks/train_llada2_common.py (run_llada2_training); train_llada2_bd.py / _with_dparallel.py become thin shims. - Model self-registration via models/llada2_moe/__init__.py (MODEL_CONFIG_REGISTRY / MODELING_REGISTRY) instead of ModelRegistry.register_modeling_path. - EP plan fix: ParallelPlan(extra_parallel_plan={'ep': ...}). - veomni_clip_grad_norm, build_dataset, save_hf_safetensor, ops_implementation. - New nested SFT configs + NPU configs (NPU code unused on this GPU). Conflict resolved in tasks/dataset/__init__.py: kept the fork's nemotron exports. PR inclusionAI#22 was validated upstream on Ascend 910B2 (loss/logits/grad parity=0.0 vs legacy, incl. real LLaDA2.0-mini weights). NOT runtime-verified in this env (Windows/py3.14, VeOmni not installable). Merged files pass py_compile. Does NOT touch the fork's long-context files; train_longctx.py + configs/longctx still use the OLD API and are migrated in the next commit.
PR inclusionAI#22 migrated the SFT entrypoints but not the fork's long-context files. This ports them to VeOmni 8ca09d7 without duplicating the training loop: - train_llada2_common.py: add backwards-compatible transform_builder / dataset_builder hooks to run_llada2_training so other entrypoints can reuse the validated loop. - train_longctx.py: rewritten from a ~550-line copy of the old flat-API loop into a thin extension using those hooks (text transform + Nemotron streaming). block_diffusion_mode=false reproduces the prior non-BD behaviour exactly. - configs/longctx/{64k,128k,256k}: migrated flat -> nested schema; field names (cp_size, offload_config.enable_activation, optimizer.lr_min, checkpoint.save_steps) verified against veomni/arguments/arguments_types.py. Verified: py_compile on all changed .py; all YAML parses; VeOmni symbols used by moe_convertor.py / dataset.py confirmed still exported at 8ca09d7. NOT runtime-verified: VeOmni not importable on Windows/py3.14. train_longctx.py and configs/longctx are my ports, NOT part of PR inclusionAI#22's NPU validation.
Transcribes VeOmni 8ca09d7's authoritative dependency constraints into a root requirements.txt for the pip install path (uv --extra gpu remains recommended): - Python >=3.11,<3.13 (this box's 3.14 is out of range — documented blocker). - torch 2.11.0+cu130, transformers 5.9.0, datasets<=2.21.0, tiktoken, einops, etc. - accelerate intentionally absent (VeOmni is FSDP2-native; nothing imports it). Findings recorded in MIGRATION_NOTES: - transformers 4.52.3 ceiling is inference-side (bundled checkpoint code); training uses the repo's own modeling file, import-validated under 5.9.0 by PR inclusionAI#22. - torch FSDP / device-mesh deprecation surface is inside VeOmni (already current), not dFactory's own code; parallel_plan._tensor import matches VeOmni convention. Verified: extend_rope_context.py --help runs (pure stdlib). Everything torch/veomni dependent is NOT runtime-verifiable here; declared explicitly.
4a — checkpoint loading (config diffs verified on real downloaded configs):
- Diffed config.json across 2.0-mini/2.1-mini/2.1-flash/2.2-flash. Shared MoE tensor
layout (256 experts, top-8, vocab 157184); 2.2 adds expert_capacity=48, block_size=32,
use_qk_norm=true, mpe=131072, rope_theta=3e6.
- configuration_llada2_moe.py: expert_capacity/block_size now explicit params, not **kwargs.
- compat.py (new, dependency-free): validate_llada2_config() fails loudly on non-LLaDA2 /
inconsistent configs and warns loudly on 2.2 block routing. Verified on all 4 real configs
+ negative cases.
- moe_convertor.py: merge/split now importable without veomni (lazy import); validation wired in.
- New model_configs/{llada2_1_mini,llada2_1_flash,llada2_2_flash} + matching configs/sft/*.
4b — long-context ceiling: documented (per-doc padding, no packing; block-diffusion mask is
O((2L)^2) and unusable long — loud guard added; longctx uses full attention). Real ceiling
depends on VeOmni CP + memory, unverifiable here.
4c — docs/UNIMPLEMENTED.md: Levenshtein editing (§3.1), block routing (top-48/256 per block),
L-EBPO (§3.2) — grounded in paper + downloaded 2.2 modeling code. editing.py stubs raise
NotImplementedError (verified), not wired into training. None implemented (unpublished, out of scope).
Verified here (stdlib): validator on real+new configs, negative cases, YAML/JSON parse, py_compile,
stubs raise. NOT verified: real weight load / conversion (needs torch+transformers+veomni+weights).
Tier 1 (VeOmni-independent, runs on CPU with only torch+pytest — RAN HERE, 11 passed):
- test_moe_convertor_roundtrip.py: synthetic llada2_moe separate-expert state dict ->
moe_merge -> split_moe_experts is lossless (exact torch.equal); non-expert tensors
pass through; corrupted expert dim rejected. This is the brief's required gate.
- test_compat_validation.py: validate_llada2_config accepts 2.0/2.1/2.2 shapes, flags 2.2
block routing, rejects non-LLaDA2/inconsistent configs.
Tier 2 (test_smoke_training.py): tiny eager-MoE model, one step per objective, checkpoint
round-trip; importorskip('veomni') so it's the per-bump gate on a real stack (skips here).
Tier 3 (test_integration_real_checkpoints.py): real-checkpoint merge/split, opt-in via
LLADA2_INTEGRATION=1 + LLADA2_CKPT (multi-GB/GPU; skipped by default).
Supporting change: moe_convertor.py now imports safetensors/transformers/tqdm/veomni lazily
so merge/split are importable with only torch (enables Tier 1).
CI: .github/workflows/tests.yml runs Tier 1 on Python 3.12 + CPU torch.
Verified: pytest tests/ -> 11 passed, 2 skipped on py3.14 + torch 2.13.0+cpu.
…(Phase 6) - README: per-model support table with confidence levels (2.0/2.1 High/Medium; 2.2 low-fidelity + UNIMPLEMENTED pointer), verification-honesty callout, correct env (Py3.11-3.12, VeOmni 8ca09d7, torch 2.11+cu130, transformers 5.9.0), nested-config quickstart, positional-YAML launch, long-context + tests + known-limitations sections. - MIGRATION_NOTES: Phase 6 entry + consolidated verified/unverified/broken breakdown + branch map. - Fix launch command: config YAML is a positional arg to parse_args (nargs='?'), not --config; corrected in train_longctx.py and configs/longctx/*. Verified: pytest tests/ still 11 passed / 2 skipped; py_compile; YAML parses.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
The _compute_llada2_loss function used .clamp_min(1) on the masked-token count which silently masked the case when a micro-batch had no valid supervised positions (all labels == -100). This could hide data quality issues or bugs in the data pipeline. The fix separates the denominator calculation from the clamping and adds a warning log when the denominator is zero, surfacing the empty-supervision condition while maintaining numerical stability.
The main() function was referencing args.mode which was only defined in the __main__ block. Added mode as a parameter to main() and updated the call site to pass args.mode.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Zero routing fields misclassified
- Changed the is_block_routing check to require expert_capacity and block_size to be positive values (not just non-None), preventing zero values from being incorrectly classified as LLaDA2.2 block routing.
You can send follow-ups to the cloud agent here.
Reviewed by Cursor Bugbot for commit 737e658. Configure here.
The is_block_routing check now requires expert_capacity and block_size to be positive values, not just non-None. This prevents configs with zero values (malformed or placeholder) from being incorrectly classified as LLaDA2.2 block routing checkpoints.
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
3 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.
Modernizes this fork to a current, correct, working-where-verifiable state. Full detail in
MIGRATION_NOTES.md; capability gaps indocs/UNIMPLEMENTED.md.What this does (atomic commits, one per phase)
600fe6d→8ca09d7by adopting upstream PR [codex] Migrate dFactory to latest VeOmni with NPU support inclusionAI/dFactory#22 (validated on Ascend NPU), then porttrain_longctx.py+configs/longctx/*to the new nested VeOmni API.requirements.txt(Python 3.11–3.12, torch 2.11.0+cu130, transformers 5.9.0).expert_capacity,block_size), a loud architecture validator (models/llada2_moe/compat.py),moe_convertorvalidation, new model-config dirs + SFT configs,docs/UNIMPLEMENTED.md, andNotImplementedErrorstubs for the unpublished 2.2 mechanisms.pytest tests/→ 11 passed on CPU).Verification honesty
py_compile, stub behaviour.train_longctxport, real-checkpoint conversion,requirements.txtinstall, and the Tier-2 smoke test. PR [codex] Migrate dFactory to latest VeOmni with NPU support inclusionAI/dFactory#22's VeOmni migration was hardware-validated upstream, but the long-context port here was not.Highest-value next step: run
pytest tests/(incl. Tier 2) on Linux + Py3.12 + CUDA 13 withuv sync --extra gputo flip the largest 🟡 block to ✅.Note
High Risk
Large training-path refactor (VeOmni API, shared loop, model MoE/registration) and broad config changes affect how models load and train; E2E SFT was not re-run in the authoring environment, so regressions are possible until Tier-2 smoke tests run on Linux+GPU.
Overview
This PR modernizes the fork around a VeOmni submodule bump and clearer LLaDA2 coverage, with honesty about what is verified vs not.
Training stack: SFT entrypoints (
train_llada2_bd.py, parallel-decoding variant, long-context) are thinned to call sharedrun_llada2_trainingintrain_llada2_common.py, aligned with VeOmni’s nested args (ops_implementation,accelerator,checkpoint, etc.). SFT and long-context YAMLs are migrated accordingly (e.g.fused_triton,cp_size). The model package self-registersllada2_moe_veomniwith VeOmni; MoE forward gains ops-config patching, eager stacked-expert fallback, and expert-parallel plan usesextra_parallel_plan.LLaDA2.1 / 2.2: New model-config dirs and SFT presets;
configuration_llada2_moeexposesexpert_capacity/block_size.validate_llada2_configincompat.pyfails on bad architectures and warns on 2.2 block routing.moe_convertorvalidates before merge/split and splits core merge logic so it imports with torch only (VeOmni lazy inmain()).docs/UNIMPLEMENTED.mdandediting.pystubs document that 2.2 block routing / Levenshtein editing / L-EBPO are not implemented.Quality & docs: Tiered
tests/plus.github/workflows/tests.yml(CPU pytest, no submodule) for merge/split losslessness and config validation; rootrequirements.txt, rewrittenREADME.md, andMIGRATION_NOTES.md..gitignoreadds.scratch_configs/.Reviewed by Cursor Bugbot for commit 16893b1. Bugbot is set up for automated code reviews on this repo. Configure here.