Repository navigation
Conversation
github-actions
Bot
added
checkpoint
ci
docker
documentation
dyedd
marked this pull request as draft
September 16, 2026 06:31
dyedd
force-pushed
the
refactor/new-structure
branch
from
September 21, 2026 08:41
cc474c6 to
b614ce7
Compare
Collaborator
|
engine 和 trainer 概念有点混用了，未来不是很容易扩展新的 train方式支持；你的调整方案，需要明确下每个目录的职责是否边界清楚，以及后续扩展新的功能是否容易； |
Squash of the layout refactor (engines split, data and models by family, shared code lifted to the top level, embodied preprocessing grouped under tools/embodied_data_preprocess) rebased onto origin/master as one linear commit.
dyedd
force-pushed
the
refactor/new-structure
branch
from
October 10, 2026 09:23
276589f to
4e21139
Compare
Place MCore and Torch runners next to the engine that owns their lifecycle, and retry embodied data imports after a failed optional dependency load.
dyedd
force-pushed
the
refactor/new-structure
branch
from
October 13, 2026 06:31
4e21139 to
8b96edd
Compare
- Rename the MCore engine to Megatron and move training orchestration into dedicated runners with direct training-method dispatch. - Pass model configuration and preprocessing dependencies explicitly, consolidate heterogeneous DP state, and reuse upstream Megatron helpers. - Separate Torch FP8 quantization from distributed setup and reorganize chat-template plugins and model-specific utilities. - Move single-model offline tools into examples and share Kimi K2.x NVFP4 conversion utilities across model variants. - Update launch scripts, dependency checks, tests, and bilingual guides for the revised package layout and interfaces.
- Move Megatron task and data arguments into training, keeping backend options in the engine and registering both on the same parser. - Move Torch training arguments, CLI parsing, learning-rate schedules, and profiling into training, and update their callers. - Remove unused CLI options, config fields, helpers, and write-only state across models, data pipelines, engines, and evaluation. - Fix the DreamZero evaluation action queue initialization and remove unused GR00T timing work while retaining replay profiling logs. - Update model configs, launch scripts, tests, documentation, and argument ownership rules for the revised interfaces.
- Split Megatron and Torch training packages and registries.\n- Move Torch optimizations under accelerations and extract shared helpers.\n- Update imports, docs, tests, and distributed utilities for the new layout.
- loongforge/accelerations -> loongforge/extensions - loongforge/models/common -> loongforge/models/base - update imports, registry strings, LoRA targets, ruff per-file-ignores
… explicitly - move CP tensor helpers into models/base/context_parallel.py - resolve dtype and RNG reset in training instead of models - drop the data -> engines tokenizer assertion - make chat_templates a leaf with a local tokenizer Protocol
- rewrite layering checks for models/data/engines/extensions boundaries - register models -> engines debt explicitly - add pyflakes and test_undefined_names to workflow lint - drop one-off scripts and rename the checkpoint converter driver
- single-ownership wording instead of "shared common layers" - data may import declarative model modules, never implementations - engines use models/base; models -> engines is registered debt - rename extensions and document the ChunkPipe queue contract
dyedd
marked this pull request as ready for review
October 20, 2026 13:05
dyedd
changed the title
- restore `--pretrained-checkpoint` and `--resume` in `_CheckpointIOArgs` - keep `training/torch/runner.py` argument access valid - allow Torch runs and cosmos3 argument parsing to use these options
- remove the unused `nullcontext` import - rename the ambiguous loop variable - keep whole-file Ruff checks clean for files in the PR diff
…nced code - Move Megatron run, freeze/init, EMA, recipe, and log-tensor arguments from engines/megatron/arguments.py to training/megatron/args.py, leaving only backend-execution arguments in the engine. Both providers still register on the same parser, so the CLI and the parsed namespace are unchanged. - Move Torch --pretrained-checkpoint and --resume into TrainingArgs; apply determinism, cuDNN autotune, and TF32 switches in engines/torch/initialize.py; move the switch-conflict and DMuon bucket-size checks into validate_engine_args so they fail at parse time. - Move the llm-inspector add_hooks helper to the training runner, which owns the log-tensor arguments that enable it. - Drop the local --padded-vocab-size definition; upstream Megatron provides it. - Fix GR00T-N1.7 reading an undeclared pretrained_checkpoint field, wall-oss fallback imports that could never succeed, the llava_onevision_1.5 config path in tests and example scripts, and stale DSA kernel comments. - Delete unreferenced modules and helpers: local_attention, canonical_lora, wan_ulysses, three orchestrator CLIs, LoRAMerge, forall, the ernie download helpers, an unused vla_mixin conversion, the cosmos3 demo block, and four unused YAML configs. - Register kimi-k2.6-bf16 in the model catalog. - Document the parameter ownership rule in AGENTS.md and refresh the README trees, tutorials, submit-pr module mapping, and pyproject metadata to the current layout.
…r model `examples/` had drifted away from its convention: a wildcard grouping directory sat next to real model directories, model directories mixed flat launch scripts with offline helpers, the same script lived at a different depth depending on the model, and DreamZero's data-preparation implementation had been collapsed out of `tools/` into `examples/`. Restore the split with `tools/` and give every model the same layout. DreamZero preprocessing back to tools/ - move prepare_dataset.py, precompute_features.py, validate_precomputed_feature_artifact.py and cache_precompute/ to tools/dreamzero_data_preprocess/, keeping the current (post-rename) import paths, so tools/ holds the CLIs and examples/ keeps thin wrappers - fix the config path built by precompute_dreamzero_cache.sh: the earlier refactor replaced configs/models/embodied/ with configs/models/ but the DreamZero YAMLs live in configs/models/dreamzero/ - split the README accordingly and refresh the tools table One directory per model - split kimi_k2.x/ into kimi_k2.5/ and kimi_k2.6/; the version-agnostic NVFP4 PTQ driver stays single-copy under the family root kimi_k2/checkpoint_convert/nvfp4_ptq/ - repoint COMMON_NVFP4_DIR in both wrappers and the two printed post-conversion commands One phase layout per model - every model now uses pretrain/, finetuning/, checkpoint_convert/, eval/ and preprocess/, with dataset fixtures under assets/ and only README.md at the model root - move the flat finetune/sft scripts of cosmos3, dreamzero, fastwam, groot_n1_6/7, lingbot_va, pi05, wall_oss_0_5, xvla and internvl2.5 into finetuning/; move offline helpers into preprocess/ keeping each .py next to its caller; give qwen_image and wan the pretrain/ + preprocess/ + checkpoint_convert/ split they were missing - align examples_xpu/deepseek_v3.1 with the same rule - scripts that derive the repo root from SCRIPT_DIR sat one level shallower before, so bump them to ../../.. and repoint intra-script references; the WAN converters no longer depend on the caller's cwd - make pretrain_wan2.2 use a relative DATASET_PATH like its wan2.1 sibling instead of a stale absolute path into the repo - fix comments that pointed at a non-existent pi05 sft script and at the pre-rename FastWAM script name Keep it from drifting again - add tests in tests/test_layout_paths.py: every examples*/ directory names a model under configs/models/ (or a version of one), a model directory holds only phase directories/assets/README, and phase scripts resolve the repo root three levels up - document the convention in support_new_model (en/zh) Update tests/embodied/config/scripts.yaml, the README model list, the support_model tables and every doc/config comment that referenced a moved script.
- models/ no longer imports engines/: the training method injects MultimodalRuntime and HeteroStepContext - ChunkPipe metadata travels with the batch instead of a shared queue - padding rules live in training/megatron/data_padding.py; the data layer only applies them - diffusion CP broadcast uses the generic broadcast primitives from the method side - the MTP logging scale is computed in sft_llm; the model no longer reads DP size or microbatch counts - DynamicRotaryEmbedding moves to models/base, with its three defects fixed - LoRA shape has a single source (the YAML); --patch-size, --lora-alpha and --lora-dim are removed - model YAML keys that shadow LoongForge CLI options are scanned by a test - InternVL2.5 YAMLs carry the HF rope scaling factor
…xtensions - Move the Wall-OSS-0.5 FSDP2 + DMuon wrap out of the model into extensions/torch/wall_oss_0_5; wrap_model now takes a dmuon_shard callback instead of probing convert_to_fsdp, and grad_sync_ctx suppresses DMuon gradient sync inside the engine - Resolve execution facts into DataRuntime in the training method; collators, samplers, transforms and dataset builders read it instead of backend arguments or DistributedContext, which data and models no longer import - Dispatch --dataset-format through DATASET_FORMATS so the shared dataloader no longer imports the X-VLA HDF5 builder - Remove the cuDNN writes from sft_vlda; the cookbook script opts into --deterministic-mode so the engine owns determinism - Extend test_data_layers with an AST check for engine-argument reads, TYPE_CHECKING import bans, and a DATASET_FORMATS registry check - Sync AGENTS.md and the Chinese support/eval docs --optimizer dmuon for Wall-OSS-0.5 now needs --trainer-type WallOss05Trainer.
…ckage table - Rename `CustomModelFamilies` to `DiffusionModelFamilies`; the entries are diffusion models, not custom vision models - Reference family constants in `FAMILY_GROUP` instead of string literals - Replace the duplicated embodied `model_type -> package` tables with `MODEL_PACKAGES` in `models/embodied/registry.py`; `data/embodied/registry.py` reads it and keeps only the data-only `PLAIN_MODULES` - Move `torch.nn` to `TYPE_CHECKING` so the data registry stays importable without torch - Drop the two-table consistency test and update the registry tests and AGENTS.md - Align the EN/ZH "Support New Model" guides: Wan is a diffusion model, the YAML examples carry the required `model_type`, and a new section lists the registrations a new architecture needs Verified: 39 unittest cases pass in the three registry test modules, ruff clean, no `MODEL_MODULES` or `CustomModelFamilies` left, doc link check passes. Not verified: pytest and runtime `get_collator` resolution, this machine has no torch.
- Move each family's DataLoader entry into data/<family>/dataloader.py and pass frozen data configs plus DP facts instead of Megatron args or global distributed state. - Move embodied transform and collator pipelines into <model>/transforms/, unify sampler modules as samplers.py, and rename the GR00T-N1.6 data processor. - Move Eagle3 HF processors into the groot_n1_6 model package and keep its package __init__ lazy. - Move full-hetero-DP encoder iterator assembly into engines/megatron/distributed/hetero_dp.py; methods now only prepare datasets and loaders. - Read dataloader state in engines/megatron/checkpointing.py and restore streaming SFT position and diffusion consumed samples. - Update layering tests, AGENTS.md, and docs for the new ownership rules.
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.
This PR aligns
LoongForgewith the repository's ownership and naming rules. It separates training lifecycle code from engine execution, organizes model and data code by responsibility, and makesloongforge/train.pyselect the training runner throughloongforge/models/catalog.py. The in-repo launch scripts, configs, tests, and documentation are updated to the new layout.Change Summary
What changed?
loongforge/training/{megatron,torch}/owns runners, training lifecycle, recipe/task/data arguments, method registries, learning-rate schedules, logging, profiling, and per-model training methods. The two training paths do not import each other.loongforge/engines/{megatron,torch}/owns parser and engine arguments, initialization, parallelism, model wrapping, train steps, optimizer integration, and checkpoint formats.Megatronglobal state now lives underloongforge/engines/megatron/.loongforge/models/organizes models by capability;loongforge/data/organizes pipelines by data family. Embodied evaluation moves toloongforge/evaluation/embodied/.loongforge/extensions/<engine>/<model>/owns training extensions that must know both a concrete model and an engine.loongforge/train.pyis the single training entrypoint. It uses--model-nameto select the catalog'smegatronortorchrunner. TheLoongForgeconsole script resolves to the same entrypoint.tools/remains a repository-level home for shared offline utilities. Layering tests reject cross-layer imports and stale dotted paths.Which issue or requirement does this address?
There is no separate public issue linked to this PR. The requirement comes from the ownership, dependency, and naming rules in
AGENTS.md; the PR title followsCONTRIBUTING.md.Which modules are affected?
Training dispatch and runners, engine execution, model and data organization, embodied evaluation, training extensions, configuration paths, examples, documentation, packaging, CI, and tests.
Validation
python3 -m unittest tests.test_data_layers tests.test_import_targets -q: 25 tests passed, 1 skipped..github/scripts/pr_title.js.loongforge/embodied/train.pypath require migration.CI, GPU, and Image Impact
pyflakesis added to the test extra and runs inworkflow-lint.llm_vlm/embodied.Reviewer Notes
Direct imports of moved
LoongForgemodules and the oldloongforge/embodied/train.pylaunch path are not kept through a compatibility layer. In-repo consumers are updated toloongforge/train.py, theLoongForgeconsole script, and the current module paths. External scripts that use old deep imports or launch paths must migrate.The static boundary and import-path checks pass. GPU and XPU end-to-end training were not run as part of this description update. Remaining
models -> enginesdependencies are recorded explicitly as technical debt intests/test_data_layers.py.