feat: add Bailing V3 SWE SFT support

areal/api/cli_args.py

@@ -936,11 +936,12 @@ class MegatronEngineConfig:
 
     # MoE
     moe_router_dtype: str | None = "fp32"
-    moe_shared_expert_overlap: bool = field(
-        default=False,
+    moe_shared_expert_overlap: bool | None = field(
+        default=None,
         metadata={
             "help": "Enable overlapping between shared expert computations and dispatcher communications. "
-            "Without this, the shared experts execute after the routed experts."
+            "Without this, the shared experts execute after the routed experts. "
+            "None keeps the model bridge's own default."
         },
     )
     moe_enable_deepep: bool = False
@@ -961,13 +962,14 @@ class MegatronEngineConfig:
             "Requires TransformerEngine >= 2.7.0.",
         },
     )
-    moe_router_bias_update_rate: float = field(
-        default=0.0,
+    moe_router_bias_update_rate: float | None = field(
+        default=None,
         metadata={
             "help": "Update rate for auxiliary-loss-free MoE load balancing "
             "(DeepSeek V3 style). Controls how fast expert_bias adjusts. "
-            "Default 0.0 disables bias updates; set a positive value such as "
-            "1e-3 to enable.",
+            "None keeps the model bridge's own default (AReaL bridges "
+            "disable it or derive it from the checkpoint). Set 0.0 to "
+            "disable explicitly; 1e-3 matches DeepSeek V3.",
         },
     )
     moe_z_loss_coeff: float | None = field(

areal/dataset/__init__.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import re
 from typing import TYPE_CHECKING, Optional
 
 from areal.api.cli_args import _DatasetConfig
@@ -19,10 +20,17 @@
     "virl39k",
     "hh-rlhf",
     "torl_data",
+    "swe_sft",
 ]
 
 logger = logging.getLogger("Dataset")
 
+# Matches "swe" only as a path token delimited by /, _, -, or . (e.g.
+# "swe_data/", "swe-bench", "my_swe.jsonl") so that paths merely containing
+# the trigram (e.g. "answer_sft", "/home/swetha/") fall through to the
+# generic load-from-disk fallback instead of the SWE trajectory pipeline.
+_SWE_PATH_PATTERN = re.compile(r"(?:^|[/_\-.])swe(?:[/_\-.]|$)")
+
 
 def _get_custom_dataset(
     path: str,
@@ -133,6 +141,16 @@ def _get_custom_dataset(
             max_length=max_length,
             **kwargs,
         )
+    elif _SWE_PATH_PATTERN.search(path.lower()) and type == "sft":
+        from .swe_sft import get_swe_sft_dataset
+
+        return get_swe_sft_dataset(
+            path=path,
+            split=split,
+            tokenizer=tokenizer,
+            max_length=max_length,
+            **kwargs,
+        )
     else:
         # Fallback: try loading as a generic HuggingFace dataset from disk.
         # This supports arbitrary datasets saved via dataset.save_to_disk().

areal/dataset/swe_sft/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""SWE trajectory preprocessing and SFT dataset loading.
+
+The public entry point remains :func:`get_swe_sft_dataset`. Implementation
+details are split by responsibility to keep the loader maintainable.
+"""
+
+from .pipeline import get_swe_sft_dataset
+
+__all__ = ["get_swe_sft_dataset"]