feat: add Bailing V3 SWE SFT support by dingzhiqiang · Pull Request #1598 · areal-project/AReaL

dingzhiqiang · 2026-08-11T12:16:02Z

Description

Add Bailing V3 (KDA + gated MLA + MoE) support to the Megatron engine and provide an SWE SFT dataset pipeline with reusable training entrypoints.

This ports the reviewed implementation from internal PR 2188 onto the latest public main. Internal Bailing example configs were intentionally excluded because they contain private checkpoint, storage, and cluster paths. The portable examples/swe/config.py and examples/swe/train_sft.py entrypoints are included.

Related Issue

Internal reference: inclusionAI/AReaL PR 2188

Type of Change

Checklist

I have read the Contributing Guide
Pre-commit hooks pass (pre-commit run --all-files)
Relevant tests pass; new tests added for new functionality
Documentation updated (if applicable; built with ./docs/build_all.sh)
Branch is up to date with main
Self-reviewed via /review-pr command
This PR was created by a coding agent via /create-pr
This PR is a breaking change

Breaking Change Details (if applicable):

None expected. Existing MoE config defaults now use None to preserve each model bridge default; explicit values retain override behavior.

Validation

git diff --check refs/remotes/github/main...HEAD — passed
Python compilation of changed modules and examples — passed
Targeted tests — collection blocked because the available environment lacks httpx
Full pre-commit — not run because uv and pre-commit are unavailable in this environment
Megatron expert review — no concrete correctness findings or PR blocker

Additional Context

The conflict with newer public-main chunked LM-head and NPU logic in megatron_engine.py was resolved by retaining upstream behavior and adding the Bailing V3 bridge/routing logic alongside it.

Port the reviewed internal implementation to the public main branch while preserving newer upstream engine behavior. Key changes: - Add Bailing V3 KDA, gated MLA, and MoE model support - Add SWE SFT dataset loading and cache handling - Add focused model, loader, and dataset tests Refs: inclusionAI/AReaL#2188 Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com>

Separate message processing, tokenization, pipeline orchestration, and CLI code so each concern can evolve without growing a single dataset module. Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com>

Port the follow-up SWE SFT fixes from swe-dev so Bailing V3 adaptive chat templates use structural assistant masks and consistent thinking modes. Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com>

Prevent collection-time stubs from replacing the real areal.dataset package for subsequent data-service tests.

sitabulaixizawaluduo · 2026-08-17T12:22:12Z

+            # Precision-alignment dumps (AReaL-friend tools/precision-alignment):
+            # when AREAL_DUMP_ROUTING is set, enable megatron RouterReplay
+            # recording so MoE expert indices can be captured during forward.
+            if os.environ.get("AREAL_DUMP_ROUTING", "") and hasattr(


Is this part of R3 support? It's inappropriate to appear here, suggesting it should be removed.

Removed in 6a1af81. This PR no longer enables RouterReplay or any precision-alignment routing dump behavior.

sitabulaixizawaluduo · 2026-08-17T12:22:58Z

        )

    def lr_scheduler_step(self):
+        if os.environ.get("AREAL_DUMP_ROUTING", "") or os.environ.get(


the same as before

Removed in 6a1af81. lr_scheduler_step() is restored to the normal scheduler path with no dump-related bypass.

sitabulaixizawaluduo · 2026-08-17T12:23:09Z


+            # Precision-alignment routing dump: record MoE expert indices for the
+            # first microbatch via megatron RouterReplay (enabled in initialize).
+            _routing_dump_path = os.environ.get("AREAL_DUMP_ROUTING", "")


the same as before

Removed in 6a1af81. The RouterReplay record setup in forward_step has been deleted.

sitabulaixizawaluduo · 2026-08-17T12:23:41Z

                        ),
                    )

+            if _routing_dump_path and not getattr(self, "_routing_dumped", False):


Similarly, it needs to be removed.

Removed in 6a1af81. The routing-record collection, CPU copies, and .pt writes have all been deleted.

sitabulaixizawaluduo · 2026-08-17T12:23:57Z

+
+        # Precision-alignment forward-only mode: no optimizer exists (see
+        # _create_optimizer), so skip zero_grad/step and run forward only.
+        _fwd_only = bool(


please remove it

Removed in 6a1af81. train_batch() now always follows the normal zero-grad, forward/backward, and optimizer-step path; the dump-only forward mode is gone.

sitabulaixizawaluduo · 2026-08-17T12:24:09Z

    def _create_optimizer(self, ft_spec: FinetuneSpec) -> None:
        if self.optimizer_config is None:
            return
+        if os.environ.get("AREAL_DUMP_ROUTING", "") or os.environ.get(


please remove it

Removed in 6a1af81. _create_optimizer() no longer skips optimizer creation based on dump environment variables.

sitabulaixizawaluduo · 2026-08-17T12:24:19Z

+
+            # Precision-alignment logp dump: save final per-token logprobs for
+            # the first microbatch (last PP stage only; this branch already is).
+            _logp_dump_path = os.environ.get("AREAL_DUMP_LOGP", "")


please remove it

Removed in 6a1af81. The per-token log-probability dump and associated CPU/file I/O have been deleted. The updated head passes both sglang and vLLM unit-test jobs.

Keep Bailing V3 support focused on production training behavior by removing out-of-scope routing and log-probability dump paths.

Place the CP zigzag unit tests under the root test pattern used by the GCP unit-test workflow.

daihaowz · 2026-08-21T07:16:13Z

    def _build_hf_mcore_bridge(self):
        if self.bridge_cls == "mbridge":
-            self.bridge = mbridge.AutoBridge.from_pretrained(
+            hf_config = PretrainedConfig.from_pretrained(


PretrainedConfig.from_pretrained() creates a base PreTrainedConfig. During serialization, to_dict() replaces the instance’s model_type="bailing_hybrid" with the base class’s empty model_type. As a result, regular HF checkpoints are exported with an invalid config.json, preventing AutoConfig, inference, or subsequent training from recognizing the model.
Please use AutoConfig.from_pretrained(..., trust_remote_code=True) to retain the actual config class, or always restore the original model configuration during export. A save/load round-trip test for Bailing V3 should also be added.

dingzhiqiang requested review from CormickKneey, HwVanICI, PrometheusComing, fishcrap, garrett4wade, geshi001, nuzant, rchardx and sitabulaixizawaluduo as code owners August 11, 2026 12:16

dingzhiqiang requested a review from daihaowz August 11, 2026 12:19

dingzhiqiang force-pushed the chucai.dzq/gh-bailing-v3-swe-sft branch from b51bfc5 to c3f22f9 Compare August 11, 2026 12:27

dingzhiqiang added the safe-to-test Ready to run unit-tests in a PR. label Aug 11, 2026

dingzhiqiang temporarily deployed to AReaL-unittests August 11, 2026 12:59 — with GitHub Actions Inactive

refactor(dataset): split SWE SFT loader into modules

1343abf

Separate message processing, tokenization, pipeline orchestration, and CLI code so each concern can evolve without growing a single dataset module. Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com>

dingzhiqiang added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 12, 2026

dingzhiqiang requested a review from yulangz August 12, 2026 03:31

fix(dataset): align SWE SFT with adaptive templates

723b49d

Port the follow-up SWE SFT fixes from swe-dev so Bailing V3 adaptive chat templates use structural assistant masks and consistent thinking modes. Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com>

dingzhiqiang added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 12, 2026

dingzhiqiang had a problem deploying to AReaL-unittests August 12, 2026 06:37 — with GitHub Actions Failure

test(dataset): restore modules after SWE loader tests

a29fc3c

Prevent collection-time stubs from replacing the real areal.dataset package for subsequent data-service tests.

dingzhiqiang added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 17, 2026

dingzhiqiang temporarily deployed to AReaL-unittests August 17, 2026 07:46 — with GitHub Actions Inactive

sitabulaixizawaluduo reviewed Aug 17, 2026

View reviewed changes

refactor(engine): remove precision dump hooks

6a1af81

Keep Bailing V3 support focused on production training behavior by removing out-of-scope routing and log-probability dump paths.

dingzhiqiang added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 17, 2026

dingzhiqiang temporarily deployed to AReaL-unittests August 17, 2026 12:57 — with GitHub Actions Inactive

test(models): run zigzag coverage in CI

52678d6

Place the CP zigzag unit tests under the root test pattern used by the GCP unit-test workflow.

dingzhiqiang added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 18, 2026

dingzhiqiang temporarily deployed to AReaL-unittests August 18, 2026 03:14 — with GitHub Actions Inactive

daihaowz reviewed Aug 21, 2026

View reviewed changes

Uh oh!

Conversation

dingzhiqiang commented Aug 11, 2026

Description

Related Issue

Type of Change

Checklist

Validation

Additional Context

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants