feat: add Qwen Flash Next training and compact 256K recipes

areal/api/cli_args.py

@@ -260,6 +260,13 @@ class GenerationHyperparameters:
             )
         },
     )
+    reward_normalization_use_std: bool = field(
+        default=True,
+        metadata={
+            "help": "Divide grouped rollout rewards by their standard deviation. "
+            "Set False to subtract only the mean when reward_normalization is enabled."
+        },
+    )
     drop_incomplete_group: bool = field(
         default=False,
         metadata={
@@ -325,6 +332,7 @@ def to_openai_agents_model_settings_dict(
     # OpenAI client kwargs even when users enable them.
     _WORKFLOW_ONLY_ARGS: ClassVar[set[str]] = {
         "reward_normalization",
+        "reward_normalization_use_std",
         "drop_incomplete_group",
     }
 
@@ -929,6 +937,15 @@ class MegatronEngineConfig:
             )
         },
     )
+    language_model_only: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "For multimodal ModelScope bridges, build only the language model. "
+                "Set True for text-only SWE inputs and False for vision inputs."
+            )
+        },
+    )
     # Don't use MegatronOptimizerConfig here because OmegaConf
     # does not recognize the annotation "torch.dtype"
     overlap_param_gather_with_optimizer_step: bool = False
@@ -1074,8 +1091,8 @@ class MegatronEngineConfig:
     bridge_type: str = field(
         default="mbridge",
         metadata={
-            "help": "Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'.",
-            "choices": ["mbridge", "megatron-bridge"],
+            "help": "Bridge backend for MegatronEngine. Choices: 'mbridge', 'megatron-bridge', or 'mcore-bridge'.",
+            "choices": ["mbridge", "megatron-bridge", "mcore-bridge"],
         },
     )
 
@@ -1089,7 +1106,7 @@ class MegatronEngineConfig:
     use_bridge_for_update_weights: bool = field(
         default=False,
         metadata={
-            "help": "When True and bridge_type='megatron-bridge', delegate live "
+            "help": "When True and bridge_type is 'megatron-bridge' or 'mcore-bridge', delegate live "
             "weight sync to bridge.export_hf_weights instead of the hand-rolled "
             "convert_to_hf registry. Required for models without a registry entry "
             "(e.g. Qwen3.5). FP8 paths fall back to the registry automatically.",
@@ -1114,6 +1131,17 @@ class MegatronEngineConfig:
         },
     )
 
+    freeze_ple_table: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Freeze the Qwen4-Exp PLE/N-gram table while keeping the PLE "
+                "projection, norm, convolution, and token embedding trainable. "
+                "Set False to train and export the full PLE table."
+            )
+        },
+    )
+
     enable_mtp_training: bool = field(
         default=False,
         metadata={

areal/api/engine_api.py

@@ -203,6 +203,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Submit a batch of requests and wait for results.
 
@@ -247,6 +248,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Prepare a batch of data for training from a dataloader.
 
@@ -751,6 +753,7 @@ def submit(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> int:
         """Submit a request to the inference engine and return immediately.
 
@@ -860,6 +863,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Submit a batch of requests to the inference engine and wait for the results.
 
@@ -911,6 +915,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Asynchronously submit and wait until a full batch is ready with controlled staleness.
 

areal/dataset/gsm8k.py

@@ -1,15 +1,32 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+from pathlib import Path
+
 from datasets import load_dataset
 
 
+def _rank_local_cache_dir(cache_dir: str | None) -> str | None:
+    if cache_dir is None:
+        return None
+    return str(Path(cache_dir) / f"rank-{os.getenv('RANK', '0')}")
+
+
 def get_gsm8k_sft_dataset(
     path: str,
     split: str,
     tokenizer,
     max_length: int | None = None,
+    cache_dir: str | None = None,
+    **_: object,
 ):
-    dataset = load_dataset(path=path, name="main", split=split)
+    dataset = load_dataset(
+        path=path,
+        name="main",
+        split=split,
+        cache_dir=_rank_local_cache_dir(cache_dir),
+        keep_in_memory=True,
+    )
 
     def process(sample):
         seq_token = tokenizer.encode(
@@ -19,11 +36,15 @@ def process(sample):
         loss_mask = [0] * len(prompt_token) + [1] * (len(seq_token) - len(prompt_token))
         return {"input_ids": seq_token, "loss_mask": loss_mask}
 
-    dataset = dataset.map(process).remove_columns(["question", "answer"])
+    dataset = dataset.map(process, keep_in_memory=True).remove_columns(
+        ["question", "answer"]
+    )
 
     if max_length is not None:
         # Filter out sequences longer than max_length
-        dataset = dataset.filter(lambda x: len(x["input_ids"]) <= max_length)
+        dataset = dataset.filter(
+            lambda x: len(x["input_ids"]) <= max_length, keep_in_memory=True
+        )
 
     return dataset
 
@@ -33,8 +54,16 @@ def get_gsm8k_rl_dataset(
     split: str,
     tokenizer,
     max_length: int | None = None,
+    cache_dir: str | None = None,
+    **_: object,
 ):
-    dataset = load_dataset(path=path, name="main", split=split)
+    dataset = load_dataset(
+        path=path,
+        name="main",
+        split=split,
+        cache_dir=_rank_local_cache_dir(cache_dir),
+        keep_in_memory=True,
+    )
 
     def process(sample):
         messages = [

areal/engine/awex/colocate_reader.py

@@ -28,11 +28,13 @@
 
 from __future__ import annotations
 
+import os
 from typing import Any
 
 import torch
 
 from areal.engine.awex.memory_saver import patch_tms_hook_mode
+from areal.engine.awex.metadata import serialize_metadata_gc
 
 # Must run before any awex import: awex.models.registry auto-imports model
 # modules at module load, and the BailingMoe module's transitive megatron import
@@ -60,13 +62,25 @@ def __init__(self, *args, model: torch.nn.Module, **kwargs):
         self._model_device = device
         super().__init__(*args, model=model, **kwargs)
 
+    def _init_reader_in_colocate_mode(self):
+        from areal.engine.awex.colocate_transport import (
+            _BoundedMemoryNcclColocateStreamBatchTransport,
+        )
+
+        super()._init_reader_in_colocate_mode()
+        self.colocate_transport = _BoundedMemoryNcclColocateStreamBatchTransport(
+            self.transfer_rank, self.infer_world_size
+        )
+
     def _set_device(self) -> None:
         # TODO(agent): This adapter is CUDA-only; physical ids remain metadata
         # identities and must not be used as CUDA_VISIBLE_DEVICES indices.
         torch.cuda.set_device(self._model_device)
         self.barrier_device = self._model_device.index
         self.backend = "nccl"
-        self.ready_tensor = torch.tensor(1, device=self._model_device)
+        self.ready_tensor = torch.tensor(
+            1, dtype=torch.int64, device=self._model_device
+        )
         logger.info(
             "Bound AWEX reader rank %s to model device %s",
             self.transfer_rank,
@@ -228,6 +242,7 @@ def _select_rank0(raw_meta_list):
     def get_model_arch_name(self) -> str:
         return self._model_arch_name
 
+    @serialize_metadata_gc
     def get_parameters_meta(self):
         return self._build_params_meta()
 
@@ -267,6 +282,9 @@ def __init__(self, scheduler: Any):
         self._infer_params_meta = None
         self._infer_conf: dict | None = None
         self._initialized = False
+        self._qwen4_frozen_contract = None
+        self._qwen4_frozen_binder = None
+        self._qwen4_parameter_identity = None
 
     # ── model / context helpers ───────────────────────────────────────
 
@@ -289,16 +307,10 @@ def _build_model_context(self) -> dict[str, Any]:
         pp_size = int(getattr(server_args, "pp_size", 1))
         dp_size = int(getattr(server_args, "dp_size", 1))
 
+        from areal.engine.sglang_fork_contract import resolve_scheduler_parallel_attr
+
         def rank_attr(name: str) -> int | None:
-            for obj in (
-                scheduler,
-                getattr(scheduler, "ps", None),
-                getattr(scheduler, "tp_worker", None),
-            ):
-                value = getattr(obj, name, None) if obj is not None else None
-                if value is not None:
-                    return int(value)
-            return None
+            return resolve_scheduler_parallel_attr(scheduler, name)
 
         tp_rank = rank_attr("tp_rank")
         if tp_rank is None and self._instance_local_rank is not None:
@@ -418,12 +430,45 @@ def _build_instance_params_meta(self):
         )
         return resolver.get_parameters_meta()
 
+    def _bind_qwen4_frozen_contract(self, timeout_s: float) -> None:
+        from pathlib import Path
+
+        from sglang.srt.managers.scheduler_components import weight_updater
+
+        from areal.models.mcore.qwen4_exp_awex import register_qwen4_exp_awex
+        from areal.models.mcore.qwen4_exp_awex_binding import SglangFrozenBinder
+        from areal.models.mcore.qwen4_exp_awex_contract import load_frozen_contract
+
+        manifest = os.environ.get("QWEN_AWEX_FROZEN_CONTRACT")
+        if not manifest:
+            raise ValueError("Qwen4Exp AWEX requires QWEN_AWEX_FROZEN_CONTRACT")
+        contract = load_frozen_contract(
+            Path(manifest), Path(self._scheduler.server_args.model_path)
+        )
+        train_info = self._meta_server_client.get_object(
+            "awex_train_info", timeout=timeout_s
+        )
+        if (
+            train_info.get("train_world_size") != self._train_world_size
+            or train_info.get("qwen4_exp_frozen_contract") != contract.to_dict()
+        ):
+            raise ValueError("Training and inference frozen contracts differ")
+        binder = SglangFrozenBinder(self._get_model, contract, weight_updater)
+        register_qwen4_exp_awex(sglang_binder=binder)
+        self._qwen4_frozen_contract = contract
+        self._qwen4_frozen_binder = binder
+
     def get_weight_metadata(self):
         """Inference-side parameters_meta for ONE engine instance."""
         if self._engine_rank is None:
             raise RuntimeError(
                 "AwexColocateReader must be initialized before getting weight metadata"
             )
+        if (
+            type(self._get_model()).__name__ == "Qwen4ExpForConditionalGeneration"
+            and self._qwen4_frozen_binder is None
+        ):
+            raise RuntimeError("Qwen4Exp metadata requires matching frozen contracts")
         if self._infer_params_meta is None:
             self._infer_params_meta = self._build_instance_params_meta()
         return self._infer_params_meta
@@ -487,6 +532,8 @@ def initialize(
 
         host, port = meta_server_addr.rsplit(":", 1)
         self._meta_server_client = MetaServerClient(host, int(port))
+        if type(self._get_model()).__name__ == "Qwen4ExpForConditionalGeneration":
+            self._bind_qwen4_frozen_contract(timeout_s)
 
         # Compute single-instance parameters_meta (also reused as the native
         # reader's constructor arg later).
@@ -511,6 +558,10 @@ def initialize(
             # awex.
             "router_dtype": _get_router_dtype(awex_hf_config),
         }
+        if self._qwen4_frozen_contract is not None:
+            infer_conf["qwen4_exp_frozen_contract"] = (
+                self._qwen4_frozen_contract.to_dict()
+            )
         self._infer_conf = infer_conf
 
         # Only one rank publishes the engine-instance-wide info the writer waits
@@ -569,6 +620,11 @@ def _ensure_reader(self) -> NCCLWorkerWeightsReader:
             reader.meta_server_client, self._local_gpu_id
         )
         reader.initialize()
+        if self._qwen4_frozen_contract is not None:
+            self._qwen4_parameter_identity = tuple(
+                (name, id(parameter))
+                for name, parameter in self._get_model().named_parameters()
+            )
         self._reader = reader
         logger.info(
             "Constructed native NCCLWorkerWeightsReader (transfer_rank=%d, "
@@ -591,6 +647,16 @@ def update_weights(self, version: int) -> None:
         if not self._initialized:
             raise RuntimeError("AwexColocateReader not initialized")
         reader = self._ensure_reader()
+        if self._qwen4_frozen_contract is not None:
+            identity = tuple(
+                (name, id(parameter))
+                for name, parameter in self._get_model().named_parameters()
+            )
+            if identity != self._qwen4_parameter_identity:
+                raise RuntimeError(
+                    "Qwen4Exp Parameters changed; reconstruct the AWEX reader after recovery"
+                )
+            reader.weight_converter.refresh_frozen_contract()
         reader.update_weights(step_id=version)
         self._rebuild_derived_weights()
         logger.info("Colocate weight update completed: version=%d", version)