fix gr00t checkpoint portability (#1219)

## Summary

This PR fixes two GR00T N1.5 checkpoint portability issues.

### 1. Processor device portability

Previously, `policy_preprocessor.json` could save
`device_processor.device` from the training environment. For example, a
checkpoint trained on MUSA could contain `device: musa`, which would
fail when serving on Ascend/NPU.

This PR overrides `device_processor.device` at serve/inference load time
using the runtime device from the current config.

### 2. Base model path portability

GR00T N1.5 checkpoints require `base_model_path` to load the base model
during serve/inference. If `base_model_path` is empty, serving can fail
with a HuggingFace repo id validation error.

This PR:
- Saves local `checkpoint_dir` values as absolute `base_model_path`
values.
- Preserves HuggingFace repo ids as-is.
- Adds a clear error when `base_model_path` is empty.

## Validation

Validated on Ascend 910B:

- Ran GR00T N1.5 5-step training and saved a checkpoint.
- Verified `config.json` contains a non-empty absolute
`base_model_path`.
- Verified `base_model_path` exists in the runtime environment.
- Manually changed `policy_preprocessor.json` device to `musa` to
simulate a cross-device checkpoint.
- Served the checkpoint with runtime device `npu`.
- `/healthz` returned `OK`.
- WebSocket inference returned `actions shape: (16, 7)`.

Author: yilongma0110

flagscale/inference/inference_pi.py

@@ -14,6 +14,7 @@
 from flagscale.models.utils.constants import ACTION, OBS_STATE
 from flagscale.platforms import get_platform  # noqa: F401 must be before model imports
 from flagscale.runner.utils import logger
+from flagscale.train.processor.pipeline import get_device_override
 from flagscale.train.train_pi import make_pre_post_processors
 
 
@@ -126,7 +127,7 @@ def run_inference(config_path: str):
 
     processor_kwargs = {}
     processor_kwargs["preprocessor_overrides"] = {
-        "device_processor": {"device": engine_cfg.device},
+        **get_device_override(getattr(engine_cfg, "device", None)),
         "normalizer_processor": {
             "stats": dataset_stats,
             "features": {**policy_config.input_features},

flagscale/inference/inference_qwen_gr00t.py

@@ -10,6 +10,7 @@
 from flagscale.models.vla import TrainablePolicy
 from flagscale.platforms import get_platform  # noqa: F401 must be before model imports
 from flagscale.train.processor import PolicyProcessorPipeline
+from flagscale.train.processor.pipeline import get_device_override
 
 
 def load_image(image_path: str, size: tuple[int, int] | None = None) -> torch.Tensor:
@@ -35,11 +36,13 @@ def run_inference(config_path: str):
     generate_cfg = cfg.generate
 
     pretrained_dir = engine_cfg.model
-    model = TrainablePolicy.from_pretrained(pretrained_dir, device=engine_cfg.device)
+    runtime_device = getattr(engine_cfg, "device", None) or "cpu"
+    model = TrainablePolicy.from_pretrained(pretrained_dir, device=runtime_device)
 
     preprocessor = PolicyProcessorPipeline.from_pretrained(
         pretrained_dir,
         config_filename="policy_preprocessor.json",
+        overrides=get_device_override(runtime_device),
     )
     postprocessor = PolicyProcessorPipeline.from_pretrained(
         pretrained_dir,

flagscale/models/vla/gr00t_n1_5/configuration_gr00t_n1_5.py

@@ -18,6 +18,10 @@ class Gr00tN15Config(PreTrainedConfig):
     # Path to the base GR00T N1.5 pretrained model (local or HuggingFace hub ID)
     base_model_path: str = "nvidia/GR00T-N1.5-3B"
 
+    # Whether to initialize model weights from base_model_path.
+    # Fine-tuned checkpoints set this to False and load weights from model.safetensors.
+    load_pretrained: bool = True
+
     # Fine-tuning control flags (passed to GR00TN15.from_pretrained)
     tune_llm: bool = False
     tune_visual: bool = False

flagscale/models/vla/gr00t_n1_5/modeling_gr00t_n1_5.py

@@ -33,14 +33,23 @@
   from LeRobot, see `GrootPolicy.finetune_with_groot_runner` below.
 """
 
+import dataclasses
 import os
+from pathlib import Path
 
 import torch
+from safetensors.torch import save_file
 from torch import Tensor
 
+from flagscale.models.utils.constants import (
+    SAFETENSORS_FILE,
+    resolve_pretrained_dir,
+)
 from flagscale.models.vla.base_policy import TrainablePolicy
 from flagscale.models.vla.gr00t_n1_5.configuration_gr00t_n1_5 import Gr00tN15Config
-from flagscale.models.vla.gr00t_n1_5.gr00t_n1 import GR00TN15
+from flagscale.models.vla.gr00t_n1_5.gr00t_n1 import GR00TN15, GR00TN15Config
+
+GR00T_CONFIG_DIR = "gr00t_config"
 
 
 class Gr00tN15(TrainablePolicy):
@@ -52,13 +61,28 @@ def __init__(self, config: Gr00tN15Config):
 
         self._handle_flash_attention_compatibility()
 
-        self._groot_model = GR00TN15.from_pretrained(
-            pretrained_model_name_or_path=config.base_model_path,
-            tune_llm=config.tune_llm,
-            tune_visual=config.tune_visual,
-            tune_projector=config.tune_projector,
-            tune_diffusion_model=config.tune_diffusion_model,
-        )
+        if config.load_pretrained:
+            self._groot_model = GR00TN15.from_pretrained(
+                pretrained_model_name_or_path=config.base_model_path,
+                tune_llm=config.tune_llm,
+                tune_visual=config.tune_visual,
+                tune_projector=config.tune_projector,
+                tune_diffusion_model=config.tune_diffusion_model,
+            )
+        else:
+            gr00t_config = GR00TN15Config.from_pretrained(config.base_model_path)
+            self._groot_model = GR00TN15(
+                gr00t_config,
+                local_model_path=config.base_model_path,
+            )
+            self._groot_model.backbone.set_trainable_parameters(
+                tune_visual=config.tune_visual,
+                tune_llm=config.tune_llm,
+            )
+            self._groot_model.action_head.set_trainable_parameters(
+                tune_projector=config.tune_projector,
+                tune_diffusion_model=config.tune_diffusion_model,
+            )
 
         self._groot_model.compute_dtype = config.compute_dtype
         self._groot_model.config.compute_dtype = config.compute_dtype
@@ -68,6 +92,46 @@ def __init__(self, config: Gr00tN15Config):
         if config.output_features:
             self.output_features = config.output_features
 
+    def _save_pretrained(self, save_directory: Path, state_dict=None) -> None:
+        """Save a self-contained GR00T N1.5 checkpoint."""
+        save_directory = Path(save_directory)
+
+        gr00t_config_dir = save_directory / GR00T_CONFIG_DIR
+        gr00t_config_dir.mkdir(parents=True, exist_ok=True)
+        self._groot_model.config.save_pretrained(gr00t_config_dir)
+
+        save_config = dataclasses.replace(
+            self.config,
+            base_model_path=GR00T_CONFIG_DIR,
+            load_pretrained=False,
+        )
+        save_config._save_pretrained(save_directory)
+
+        if state_dict is not None:
+            state_dict = {k: v.clone().contiguous() for k, v in state_dict.items()}
+        else:
+            state_dict = {k: v.clone().contiguous() for k, v in self.state_dict().items()}
+        save_file(state_dict, str(save_directory / SAFETENSORS_FILE))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_path, device="cpu", *, config=None):
+        """Load a fine-tuned GR00T N1.5 checkpoint without reloading the base model."""
+        path = resolve_pretrained_dir(Path(pretrained_path), SAFETENSORS_FILE)
+        if config is None:
+            config = Gr00tN15Config.from_pretrained(path)
+
+        base_model_path = Path(config.base_model_path)
+        if not base_model_path.is_absolute():
+            base_model_path = path / base_model_path
+
+        config = dataclasses.replace(
+            config,
+            base_model_path=str(base_model_path),
+            load_pretrained=False,
+        )
+
+        return super().from_pretrained(pretrained_path, device=device, config=config)
+
     def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
         """Training forward pass.
 

flagscale/serve/run_serve_vla.py

@@ -16,6 +16,7 @@
 from flagscale.serve.processor.image_resize_processor import ImageResizeProcessorStep
 from flagscale.serve.websocket_policy_server import WebsocketPolicyServer
 from flagscale.train.processor import PolicyProcessorPipeline, ProcessorStepRegistry
+from flagscale.train.processor.pipeline import get_device_override
 
 # TODO: (yupu) to constant.py?
 TASK_KEY = "task"
@@ -136,9 +137,8 @@ def load_policy(self) -> None:
         """Load the policy model and all processor pipelines from the checkpoint."""
         t_s = time.perf_counter()
         pretrained_dir: str = self.config_engine.model
-        self.model = TrainablePolicy.from_pretrained(
-            pretrained_dir, device=self.config_engine.device
-        )
+        runtime_device = getattr(self.config_engine, "device", None) or "cpu"
+        self.model = TrainablePolicy.from_pretrained(pretrained_dir, device=runtime_device)
         self._load_processors(pretrained_dir)
         # TODO: (yupu): model.to(dtype)?
         logger.info(f"Policy model loading latency: {time.perf_counter() - t_s:.2f}s")
@@ -147,9 +147,11 @@ def _load_processors(self, pretrained_dir: str) -> None:
         """Load pre/post-processors from the checkpoint and build the serve preprocessor."""
         self.rename_map = self.config_engine.get("rename_map")
 
+        runtime_device = getattr(self.config_engine, "device", None) or "cpu"
         self.preprocessor = PolicyProcessorPipeline.from_pretrained(
             pretrained_dir,
             config_filename="policy_preprocessor.json",
+            overrides=get_device_override(runtime_device),
         )
         self.postprocessor = PolicyProcessorPipeline.from_pretrained(
             pretrained_dir,

flagscale/train/processor/pipeline.py

@@ -1455,6 +1455,13 @@ def process_complementary_data(self, complementary_data: dict[str, Any]) -> dict
 PolicyProcessorPipeline: TypeAlias = DataProcessorPipeline[TInput, TOutput]
 
 
+def get_device_override(runtime_device: str | None) -> dict[str, Any]:
+    """Return processor overrides to align checkpoint device with runtime device."""
+    if not runtime_device:
+        return {}
+    return {"device_processor": {"device": runtime_device}}
+
+
 class ObservationProcessorStep(ProcessorStep, ABC):
     """An abstract `ProcessorStep` that specifically targets the observation in a transition."""