Skip to content

refactor: split the package into engine, model, and data layers - #232

Draft
dyedd wants to merge 10 commits into
baidu-baige:masterfrom
dyedd:refactor/new-structure
Draft

refactor: split the package into engine, model, and data layers#232
dyedd wants to merge 10 commits into
baidu-baige:masterfrom
dyedd:refactor/new-structure

Conversation

@dyedd

@dyedd dyedd commented Sep 16, 2026

Copy link
Copy Markdown
Contributor

loongforge/ now separates the two training backends instead of sharing modules that only one of them needs.

  • engine/{common,mcore,native}/ owns each backend's parser, trainers, registries, distributed strategies, and checkpoint flow; contracts/, distributed/, optim/, and checkpoint/ keep only what both engines use.
  • loongforge/train.py and loongforge/embodied/train.py collapse into one entry point (python -m loongforge train --engine {mcore,native}, also the LoongForge console script), driven by configs/recipes/ and configs/engines/; the launch scripts, examples, tests, and docs follow it.
  • Model registries split per engine (models/mcore_registry.py, models/native_registry.py), and MCore's Megatron-backed globals move to engine/mcore/global_vars.py, so importing loongforge.utils no longer drags Megatron into Native runs. A regression test pins that boundary.
  • tools/ exists only at the repository root, and package-internal pytest files move to tests/native/unit/ so they stop shipping in the wheel.

Downstream launch scripts have to switch from the old file paths to -m loongforge; every in-repo script and example is already updated here.

Change Summary

  • What changed?
  • Which issue or requirement does this address?
  • Which modules are affected?

Validation

  • I ran the relevant local checks.
  • I added or updated tests where appropriate.
  • I checked compatibility with existing APIs and configurations.

CI, GPU, and Image Impact

  • This change affects dependencies, Dockerfiles, patches, or runtime environments.
  • GPU validation is required. Suggested suite: llm_vlm / embodied.
  • A local candidate image should be built during GPU validation.
  • No GPU or image validation is required.

Reviewer Notes

Describe any known risks, rollout considerations, or follow-up work.

@dyedd
refactor: split the package into engine, model, and data layers
f95a6eb 
`loongforge/` now separates the two training backends instead of sharing
modules that only one of them needs.

- `engine/{common,mcore,native}/` owns each backend's parser, trainers,
  registries, distributed strategies, and checkpoint flow; `contracts/`,
  `distributed/`, `optim/`, and `checkpoint/` keep only what both engines use.
- `loongforge/train.py` and `loongforge/embodied/train.py` collapse into one
  entry point (`python -m loongforge train --engine {mcore,native}`, also the
  `LoongForge` console script), driven by `configs/recipes/` and
  `configs/engines/`; the launch scripts, examples, tests, and docs follow it.
- Model registries split per engine (`models/mcore_registry.py`,
  `models/native_registry.py`), and MCore's Megatron-backed globals move to
  `engine/mcore/global_vars.py`, so importing `loongforge.utils` no longer
  drags Megatron into Native runs. A regression test pins that boundary.
- `tools/` exists only at the repository root, and package-internal pytest
  files move to `tests/native/unit/` so they stop shipping in the wheel.

Downstream launch scripts have to switch from the old file paths to
`-m loongforge`; every in-repo script and example is already updated here.
@github-actions github-actions Bot added checkpoint ci docker documentation Improvements or additions to documentation model training labels Sep 16, 2026
@dyedd
dyedd marked this pull request as draft September 16, 2026 06:31
@dyedd
refactor: rename the native engine to torch and the llm_vlm suite to …
5646028 
…mcore

The engine axis now reads as `mcore` (Megatron-backed) versus `torch`
(standalone) in package paths, CLI values, suite names, and CI wiring, so a
name says which training loop it refers to instead of which model family was
tested first.

- `loongforge/engine/native/` and `loongforge/checkpoint/native/` move to
  `engine/torch/` and `checkpoint/torch/`; `models/native_registry.py`,
  `NativeCheckpointMetadata`, and the `NATIVE_*` symbols and env vars follow.
- `--engine native` becomes `--engine torch`, and
  `configs/engines/native.yaml` becomes `torch.yaml`; every in-repo launch
  script, example, test, and doc is updated here.
- `tests/llm_vlm/` becomes `tests/mcore/` and the suite flag becomes
  `--suite {mcore,torch}`; `ci_command.py`, the workflows, the self-runner
  scripts, artifact names, `CI_ENABLE_MCORE`, and
  `LOONGFORGE_BASELINE_MCORE` follow.
- `docs/{source,source_zh}/native_tutorial/` becomes `torch_tutorial/`,
  including the readthedocs links in both READMEs.
- Drop `tests/torch/__init__.py`: with `tests/` on `sys.path`, pytest's
  prepend import mode resolved `import torch` to that test package.
- The torch suite's data root default moves from
  `/workspace/loongforge_native_ci` to `/workspace/loongforge_torch_ci` in
  `config/env.sh`, `cli.py`, and the example scripts, matching the mcore
  suite's `/workspace/loongforge_ci`.
@dyedd
refactor: fold the one-file packages into single modules
c253749 
Batch 1 of the structure cleanup: a package per type is not a boundary, so
`contracts/` and `distributed/` become modules and the two leaf helpers join
the modules they belong to.

- `contracts/{model,data,training,checkpoint}.py` and the re-exporting
  `__init__.py` collapse into `contracts.py`; the Cosmos/OpenMDW notice moves
  onto `Cosmos3Batch`, which is the only borrowed code there.
- `distributed/{context,collectives}.py` collapse into `distributed.py`, with
  the engine boundary (DDP/FSDP vs TP/PP/CP/VPP) stated in the docstring.
- `engine/common/__init__.py` becomes `engine/dispatch.py`: it dispatches a
  `TrainSpec`, and a package holding one `__init__.py` hid that.
- `utils/xpu_init.py` becomes `_init_xpu_plugin()` in `utils/utils.py`, and
  `tokenizer/defaults.py` joins `tokenizer/tokenizer.py`; both keep the same
  import-time behaviour.
- `models/dtype.py` stays as it is: the only merge targets would pull
  `megatron.core` into Torch-side imports or make models depend on the
  top-level runtime package.

README.md, README_zh.md, CLAUDE.md, and the structure doc follow the new paths.
@dyedd
refactor: give the model catalog one construction style
097f7d7 
`MCORE_CONFIGS` stored `{"config_path": ..., "config_name": ...}` dicts that
`get_model_spec()` reassembled into a `ModelSpec` with its own copy of the
config root, while `TORCH_CONFIGS` already held `ModelSpec` entries built by
`_torch()`. Both registries now go through `_mcore()` / `_torch()` and resolve
their YAML against the single `_CONFIGS_DIR`, so `get_model_spec()` is a
lookup plus the engine and existence checks.

All 112 entries resolve to the same engine, config file, and config classes as
before, verified by comparing `get_model_spec()` output against the previous
implementation.
@dyedd
refactor: make every package directory a regular package
816ef0f 
Seventeen directories under `loongforge/` were imported as packages but had no
`__init__.py`, so they resolved as namespace packages: two same-named
directories on `sys.path` would merge silently instead of one winning, and the
layout disagreed with the rest of the tree. Each now carries the license
header and a one-line docstring; no re-exports are added, so no module gains
side effects.

`loongforge/models/vla/xvla/tests/` is included for consistency even though
package-local test code is still a candidate to move under `tests/`.
@dyedd
refactor: name the vision tower modules by role
f96bf3b 
Batch 3 of the structure cleanup: each `models/vision/<tower>/` repeated the
tower name in every file name, so the same four roles were spelled eight
different ways.

- A tower now names its modules by role -- `config.py`, `layer_spec.py`,
  `vision_model.py`, `adapter.py`, plus `attention.py`, `transformer_layer.py`,
  `image_preprocess.py`, `merger.py`, `rope_utils.py`, and
  `patch_merger_adapter.py` where the tower needs them. The tower name lives in
  the directory, not in each file.
- `base_vision_models/` held a single module, so `BaseVisionModel` moves up to
  `models/vision/base_vision_model.py`, beside the shared ViT block.
- The `model_spec:` entries in `configs/models/` and the two FP8-for-VLM guides
  follow the new module paths. Class and function names do not change, so the
  rest of the configs stays put.

Checked by resolving every dotted `loongforge.*` reference in the repository
against the file tree and by parsing every import under `loongforge/`; no
module was imported at runtime (a CPU-only box has no torch).
@dyedd
refactor: move the xvla attention test out of the package
ed13cff 
Tests are not supposed to ship in `loongforge/` -- that is why the other
package-internal units already live under `tests/` -- and the XVLA attention
test was the last one left behind.

- `models/vla/xvla/tests/test_attention_forward.py` and its runner move to
  `tests/torch/unit/`, next to the other Torch units.
- The runner derives the repo root from its own location instead of defaulting
  to `/workspace/LoongForge`, so it works from any checkout.
- `tests/torch/README.md` lists `unit/` in the directory tree; the tree had
  never mentioned it.

The test still skips itself when CUDA is absent, and nothing else in the repo
referenced the old path.
@dyedd
fix: repair two references left stale by the layer split
e7ca635 
`f95a6eb` moved the multimodal dataloader into `engine/mcore/` and the VLM data
pipeline under `datasets/`, but two references kept the old shape.

- `engine/mcore/dataloader_provider.py` imported `print_error_handler` from
  `.base.task_encoder`, which does not resolve under `engine/mcore/`; the
  function lives in `loongforge.datasets.multimodal.base.task_encoder`.
- `models/vla/xvla/xvla_processor.py` described the training transforms as
  `loongforge.datasets.transforms.xvla.xvla_transform`, missing the `robotics`
  segment.

Found by checking every dotted `loongforge.*` reference and every import in the
repository against the file tree.
@dyedd
refactor: drop the compare_repro pass-through script
ed4ed97 
`evaluation/orchestrator/compare_repro.py` was a twelve-line `main()` wrapper
around `evaluation/metrics/repro_compare.py`, with no caller in the code, the
docs, or the scripts. The underlying module already runs on its own as
`python -m loongforge.evaluation.metrics.repro_compare`.
@dyedd
refactor: name model modules by role, not by family prefix
cc474c6 
`models/llm/qwen3/qwen3_config.py` repeated the family name the directory
already carries. The convention the vision towers adopted in the previous batch
now covers every model directory: the path names the family, the module names
the role (`config.py`, `model.py`, `layer_spec.py`, `attention.py`,
`provider.py`).

- 124 renames: the 8 vision tower directories and their 41 modules
  (`moon_vision_models/` -> `moon/`, `qwen3_5_vision_models/` -> `qwen3_5_vl/`),
  plus 83 LLM, diffusion, vision, VLA, and world modules (`deepseek_model.py` ->
  `model.py`, `kimi_k3_moe.py` -> `moe.py`, `wan_flow_match.py` ->
  `flow_match.py`, `dreamzero_provider.py` -> `provider.py`).
- `world/cosmos3/qwen3_vl.py` becomes `modeling_qwen3_vl.py` because it is a
  port of the HF file. Every other HF port (`modeling_*`,
  `model_configuration_*`, `configuration_*`) keeps its upstream name, and the
  rule is now written down in CLAUDE.md and both READMEs.
- 152 dotted `loongforge.models.*` references and 104 relative imports are
  updated, together with 67 `model_spec` entries in 37 `configs/` files and 11
  modules under `engine/`, `datasets/`, `tests/`, and `tools/`. Two DreamZero
  comments pointed at the old `wan_flow_match.py` path and follow.

Verified statically: every `loongforge.*` import in the tree resolves on disk,
`compileall` is clean, the 61 CI-helper pytest cases pass, and ruff reports the
same per-rule counts as `ed4ed97`. No GPU run: only module paths moved.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

checkpoint ci docker documentation Improvements or additions to documentation model training

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@dyedd