Add DeepSeek V4 Flash support for ART Megatron

pyproject.toml

@@ -61,7 +61,7 @@ megatron = [
     "nvidia-ml-py==13.580.82",
     "nvidia-modelopt>=0.42.0a0 ; sys_platform != 'darwin'",
     "nvidia-resiliency-ext<0.5 ; sys_platform == 'linux'",
-    "transformers==5.6.2",
+    "transformers==5.12.1",
     "ml-dtypes>=0.5.0 ; python_full_version < '3.13'",
 ]
 
@@ -194,10 +194,8 @@ name = "deep-ep"
 version = "1.2.1+9af0e0d"
 requires-dist = []
 
-# The Megatron Bridge source metadata currently requires Transformers 5.8.x,
-# but this branch is validated against Transformers 5.6.2 for Gemma 4.
 # Keep Bridge's runtime deps explicit here and let ART's megatron extra own the
-# Transformers pin.
+# Transformers pin validated by model-support handlers in this branch.
 [[tool.uv.dependency-metadata]]
 name = "megatron-bridge"
 version = "0.5.0+e1a207ac"

src/art/dev/__init__.py

@@ -9,10 +9,15 @@
     TinkerNativeArgs,
     TinkerTrainingClientArgs,
     TrainerArgs,
+    VllmRuntimeArgs,
 )
 from .openai_server import OpenAIServerConfig, ServerArgs, get_openai_server_config
 from .train import TrainConfig, TrainSFTConfig
-from .validate import is_dedicated_mode, validate_dedicated_config
+from .validate import (
+    is_dedicated_mode,
+    is_external_vllm_mode,
+    validate_dedicated_config,
+)
 
 __all__ = [
     "EngineArgs",
@@ -25,8 +30,10 @@
     "TinkerNativeArgs",
     "TinkerTrainingClientArgs",
     "TrainerArgs",
+    "VllmRuntimeArgs",
     "get_openai_server_config",
     "is_dedicated_mode",
+    "is_external_vllm_mode",
     "OpenAIServerConfig",
     "ServerArgs",
     "TrainSFTConfig",

src/art/dev/get_model_config.py

@@ -1,4 +1,7 @@
-from ..megatron.model_support import default_target_modules_for_model
+from ..megatron.model_support import (
+    default_target_modules_for_model,
+    vllm_lora_config_for_model,
+)
 from .engine import EngineArgs
 from .model import (
     PEFT_ARGS_MIGRATION_MESSAGE,
@@ -67,7 +70,11 @@ def get_model_config(
     if rollout_weights_mode == "lora" and "lora_target_modules" not in config.get(
         "engine_args", {}
     ):
-        engine_args["lora_target_modules"] = merged_lora_config["target_modules"]
+        engine_args["lora_target_modules"] = vllm_lora_config_for_model(
+            base_model,
+            dict(merged_lora_config),
+            allow_unvalidated_arch=True,
+        )["target_modules"]
     trainer_args = TrainerArgs(
         adam_beta1=0.9,
         adam_beta2=0.99,
@@ -101,4 +108,6 @@ def get_model_config(
         result["trainer_gpu_ids"] = config["trainer_gpu_ids"]
     if "inference_gpu_ids" in config:
         result["inference_gpu_ids"] = config["inference_gpu_ids"]
+    if "vllm_runtime" in config:
+        result["vllm_runtime"] = config["vllm_runtime"]
     return result

src/art/dev/model.py

@@ -6,6 +6,16 @@
 from .engine import EngineArgs
 
 RolloutWeightsMode = Literal["lora", "merged"]
+VllmRuntimeMode = Literal["managed", "external"]
+
+
+class VllmRuntimeArgs(TypedDict, total=False):
+    mode: Required[VllmRuntimeMode]
+    server_url: str
+    api_key: str | None
+    local_checkpoint_root: str | None
+    server_checkpoint_root: str | None
+    health_timeout_s: float
 
 
 # Vendored from transformers.training_args.OptimizerNames
@@ -135,6 +145,8 @@ class InternalModelConfig(TypedDict, total=False):
         chat_template_content_format: vLLM chat template content format.
         chat_template_tool_schema_format: Tool schema rendering format used for
             local training tokenization.
+        vllm_runtime: vLLM runtime location. Omit for ART-managed local runtime;
+            set mode="external" to attach to a pre-launched vLLM server.
         allow_unvalidated_arch: Permit model-support validation workflows to run
             architectures that are not yet in the supported-model registry.
     """
@@ -152,6 +164,7 @@ class InternalModelConfig(TypedDict, total=False):
     chat_template_path: str
     chat_template_content_format: str
     chat_template_tool_schema_format: Literal["default", "vllm_openai"]
+    vllm_runtime: VllmRuntimeArgs
     allow_unvalidated_arch: bool
 
 

src/art/dev/validate.py

@@ -1,11 +1,30 @@
 """Validation functions for model configuration."""
 
-from .model import InternalModelConfig, RolloutWeightsMode
+from collections.abc import Mapping
+from typing import cast
+
+from .model import InternalModelConfig, RolloutWeightsMode, VllmRuntimeMode
+
+
+def _vllm_runtime_mode(config: InternalModelConfig) -> VllmRuntimeMode:
+    runtime_config = config.get("vllm_runtime", {})
+    if not isinstance(runtime_config, Mapping):
+        raise ValueError("vllm_runtime must be a mapping")
+    mode = runtime_config.get("mode", "managed")
+    if mode in {"managed", "external"}:
+        return cast(VllmRuntimeMode, mode)
+    raise ValueError("vllm_runtime.mode must be either 'managed' or 'external'")
+
+
+def is_external_vllm_mode(config: InternalModelConfig) -> bool:
+    return _vllm_runtime_mode(config) == "external"
 
 
 def is_dedicated_mode(config: InternalModelConfig) -> bool:
     """Return True if the config specifies dedicated mode (separate training and inference GPUs)."""
-    return "trainer_gpu_ids" in config and "inference_gpu_ids" in config
+    return is_external_vllm_mode(config) or (
+        "trainer_gpu_ids" in config and "inference_gpu_ids" in config
+    )
 
 
 def _rollout_weights_mode(config: InternalModelConfig) -> RolloutWeightsMode:
@@ -24,6 +43,25 @@ def validate_dedicated_config(config: InternalModelConfig) -> None:
     has_trainer = "trainer_gpu_ids" in config
     has_inference = "inference_gpu_ids" in config
     rollout_weights_mode = _rollout_weights_mode(config)
+    external = is_external_vllm_mode(config)
+
+    if external:
+        runtime_config = config.get("vllm_runtime", {})
+        assert isinstance(runtime_config, Mapping)
+        if not runtime_config.get("server_url"):
+            raise ValueError("vllm_runtime.server_url is required for external mode")
+        if rollout_weights_mode != "lora":
+            raise ValueError(
+                "vllm_runtime.mode='external' requires rollout_weights_mode='lora'"
+            )
+        if has_trainer and not config["trainer_gpu_ids"]:
+            raise ValueError("trainer_gpu_ids must be non-empty")
+        if "fast_inference" in config.get("init_args", {}):
+            raise ValueError(
+                "fast_inference is no longer supported; ART always uses an external "
+                "vLLM runtime"
+            )
+        return
 
     if has_trainer != has_inference:
         raise ValueError(

src/art/local/backend.py

@@ -43,6 +43,10 @@
     pull_model_from_s3,
     push_model_to_s3,
 )
+from art.vllm_runtime import (
+    get_external_vllm_runtime_config,
+    openai_base_url_from_vllm_server_url,
+)
 from mp_actors import close_proxy, move_to_child_process
 
 from .. import dev
@@ -111,6 +115,16 @@ def _configured_chat_template_server_arg(
     return chat_template_path or chat_template
 
 
+def _model_support_default_chat_template(
+    base_model: str,
+    internal_config: dev.InternalModelConfig,
+) -> str | None:
+    handler = _model_support_handler(base_model, internal_config)
+    if handler is None:
+        return None
+    return handler.default_chat_template()
+
+
 def _apply_configured_chat_template(
     tokenizer: PreTrainedTokenizerBase,
     internal_config: dev.InternalModelConfig,
@@ -120,11 +134,37 @@ def _apply_configured_chat_template(
         tokenizer.chat_template = chat_template
 
 
+def _model_support_handler(
+    base_model: str,
+    internal_config: dev.InternalModelConfig,
+) -> Any | None:
+    from ..megatron.model_support.registry import (
+        UnsupportedModelArchitectureError,
+        get_model_support_handler,
+    )
+
+    try:
+        return get_model_support_handler(
+            base_model,
+            allow_unvalidated_arch=bool(
+                internal_config.get("allow_unvalidated_arch", False)
+            ),
+        )
+    except UnsupportedModelArchitectureError:
+        return None
+
+
 def _apply_configured_chat_template_server_args(
     config_dict: dict,
     internal_config: dev.InternalModelConfig,
+    *,
+    base_model: str | None = None,
 ) -> None:
     chat_template = _configured_chat_template_server_arg(internal_config)
+    if chat_template is None and base_model is not None:
+        chat_template = _model_support_default_chat_template(
+            base_model, internal_config
+        )
     if chat_template is None:
         return
     server_args = dict(config_dict.get("server_args", {}))
@@ -300,6 +340,19 @@ def _chat_template_tool_schema_format(
             self._default_chat_template_tool_schema_format,
         )
 
+    def _configure_training_tokenizer(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        *,
+        model: AnyTrainableModel,
+        internal_config: dev.InternalModelConfig,
+    ) -> PreTrainedTokenizerBase:
+        _apply_configured_chat_template(tokenizer, internal_config)
+        handler = _model_support_handler(model.base_model, internal_config)
+        if handler is None:
+            return tokenizer
+        return handler.configure_tokenizer(tokenizer, internal_config=internal_config)
+
     def __enter__(self) -> Self:
         return self
 
@@ -524,8 +577,11 @@ def _get_packed_tensors(
         internal_config = cast(dev.InternalModelConfig, model._internal_config or {})
         tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
         if tokenizer_key not in self._tokenizers:
-            tokenizer = AutoTokenizer.from_pretrained(model.base_model)
-            _apply_configured_chat_template(tokenizer, internal_config)
+            tokenizer = self._configure_training_tokenizer(
+                AutoTokenizer.from_pretrained(model.base_model),
+                model=model,
+                internal_config=internal_config,
+            )
             self._tokenizers[tokenizer_key] = tokenizer
         if model.base_model not in self._image_processors:
             try:
@@ -702,8 +758,15 @@ async def _prepare_backend_for_training(
         service = await self._get_service(model)
         host, port = await service.start_openai_server(config=resolved_config)
 
-        base_url = f"http://{host}:{port}/v1"
-        api_key = server_args.get("api_key") or "default"
+        external_runtime = get_external_vllm_runtime_config(internal_config)
+        if external_runtime is not None:
+            base_url = openai_base_url_from_vllm_server_url(external_runtime.server_url)
+            api_key = (
+                server_args.get("api_key") or external_runtime.api_key or "default"
+            )
+        else:
+            base_url = f"http://{host}:{port}/v1"
+            api_key = server_args.get("api_key") or "default"
 
         return base_url, api_key
 
@@ -1130,8 +1193,11 @@ async def _train_sft(
         internal_config = cast(dev.InternalModelConfig, model._internal_config or {})
         tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
         if tokenizer_key not in self._tokenizers:
-            tokenizer = AutoTokenizer.from_pretrained(model.base_model)
-            _apply_configured_chat_template(tokenizer, internal_config)
+            tokenizer = self._configure_training_tokenizer(
+                AutoTokenizer.from_pretrained(model.base_model),
+                model=model,
+                internal_config=internal_config,
+            )
             self._tokenizers[tokenizer_key] = tokenizer
         tokenizer = self._tokenizers[tokenizer_key]
 

src/art/megatron/dsv4/__init__.py

@@ -0,0 +1 @@
+