[ Train] [Features] add new features of Orca model in qwen_gr00t

examples/qwen_gr00t/conf/train/qwen_gr00t.yaml

@@ -19,6 +19,24 @@ system:
     # Path to a checkpoint directory to resume training from (e.g. /path/to/checkpoints/005000)
     # resume_from:
 
+  # --- Activation Checkpointing ---
+  # Modes: "none" (default), "full", "selective", "memory_budget"
+  # "full": wraps every FSDP unit with gradient checkpointing
+  # "selective": wraps every Nth layer (selective_ac_option="2" = every 2nd)
+  #              or uses per-op SAC (selective_ac_option="op")
+  # "memory_budget": compiler-driven (requires torch.compile)
+  activation_checkpoint:
+    mode: none
+    # For selective mode: "2" = every 2nd layer, "3" = every 3rd, "op" = per-op SAC
+    # selective_ac_option: "2"
+    preserve_rng_state: false
+    # For memory_budget mode only:
+    # memory_budget: 0.5
+    # Only checkpoint the action model (DiT blocks). VLM layers use dynamic
+    # sequence lengths (variable vision tokens) which breaks recomputation.
+    # checkpoint_patterns:
+    #   - "action_model\\._head\\.model\\.transformer_blocks\\.\\d+$"
+
 model:
   model_name: qwen_gr00t
   vlm:
@@ -56,6 +74,12 @@ model:
 
   prompt_template: "Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format."
 
+  # --- Chunked Cross Entropy ---
+  # 0 = disabled
+  # >0 = chunk size in tokens. Processes logits in fixed-size chunks to reduce
+  #       peak memory from the [batch*seq_len, vocab_size] float32 softmax tensor.
+  chunked_ce_tokens: 0
+
   normalization_mapping:
     VISUAL: IDENTITY
     STATE: MIN_MAX
@@ -123,6 +147,21 @@ data:
   # Path to VLM co-training data (WDS/Energon format). Leave unset to disable co-training.
   # vlm_data_path: /workspace/datasets/vlm_cotrain/
   tolerance_s: 0.0001
+
+  # --- Data Mixture ---
+  # When data_mix is set, it takes precedence over data_path.
+  # Each entry specifies a dataset path and sampling weight.
+  # Weights are relative — they get normalized internally.
+  # Optional: data_root_dir as a common prefix for relative paths.
+  #
+  # data_root_dir: /workspace/datasets/
+  # data_mix:
+  #   - path: IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot
+  #     weight: 1.0
+  #   - path: IPEC-COMMUNITY/libero_spatial_no_noops_1.0.0_lerobot
+  #     weight: 0.5
+  # balance_dataset_weights: true  # scale weights by dataset size (default: true)
+
   preprocessor:
     name: policy_preprocessor
     steps:

flagscale/models/utils/constants.py

@@ -49,6 +49,7 @@
 OPTIMIZER_STATE = "optimizer_state.safetensors"
 OPTIMIZER_PARAM_GROUPS = "optimizer_param_groups.json"
 SCHEDULER_STATE = "scheduler_state.json"
+DATALOADER_STATE = "dataloader_state.json"
 
 POLICY_PREPROCESSOR_DEFAULT_NAME = "policy_preprocessor"
 POLICY_POSTPROCESSOR_DEFAULT_NAME = "policy_postprocessor"

flagscale/models/vla/qwen_gr00t/configuration_qwen_gr00t.py

@@ -23,6 +23,10 @@ class QwenGr00tConfig(PreTrainedConfig):
 
     prompt_template: str | None = None
 
+    # Chunked cross-entropy for VLM co-training loss.
+    # 0 = disabled (use HF model's built-in CE), >0 = chunk size in tokens.
+    chunked_ce_tokens: int = 0
+
     normalization_mapping: dict[str, NormalizationMode] = field(
         default_factory=lambda: {
             "VISUAL": NormalizationMode.IDENTITY,
@@ -61,8 +65,14 @@ def from_train_config(cls, train_config: TrainConfig) -> QwenGr00tConfig:
         action_model = GR00TActionHeadConfig.from_omegaconf(model_cfg.action_model)
 
         prompt_template = getattr(model_cfg, "prompt_template", None)
+        chunked_ce_tokens = getattr(model_cfg, "chunked_ce_tokens", 0)
 
-        kwargs = dict(vlm=vlm, action_model=action_model, prompt_template=prompt_template)
+        kwargs = dict(
+            vlm=vlm,
+            action_model=action_model,
+            prompt_template=prompt_template,
+            chunked_ce_tokens=chunked_ce_tokens,
+        )
 
         raw_norm = getattr(model_cfg, "normalization_mapping", None)
         if raw_norm is not None:

flagscale/models/vla/qwen_gr00t/modeling_qwen_gr00t.py

@@ -32,6 +32,7 @@
 from flagscale.models.vla.registry import build_action_model, build_vlm
 from flagscale.models.vla.utils import get_vlm_config
 from flagscale.platforms.platform_manager import get_platform
+from flagscale.train.utils.chunked_cross_entropy import chunked_cross_entropy_loss
 
 
 class QwenGr00t(TrainablePolicy):
@@ -158,7 +159,16 @@ def forward(
 
         if vlm_batch is not None:
             with torch.autocast(get_platform().amp_device_type(), dtype=torch.bfloat16):
-                vlm_loss = self.vlm.model(**vlm_batch, return_dict=True).loss
+                if self.config.chunked_ce_tokens > 0:
+                    labels = vlm_batch.pop("labels")
+                    output_vlm = self.vlm.model(**vlm_batch, return_dict=True)
+                    vlm_loss = chunked_cross_entropy_loss(
+                        output_vlm.logits,
+                        labels,
+                        chunk_tokens=self.config.chunked_ce_tokens,
+                    )
+                else:
+                    vlm_loss = self.vlm.model(**vlm_batch, return_dict=True).loss
             result["vlm_loss"] = vlm_loss
 
         return result

flagscale/models/vla/vlm/qwenvl_backbone.py

@@ -18,7 +18,6 @@
     Qwen3VLForConditionalGeneration,
 )
 
-from flagscale.logger import logger
 from flagscale.models.vla.registry import register_vlm
 from flagscale.platforms.platform_manager import get_platform
 
@@ -91,14 +90,9 @@ def prepare_input(
         if isinstance(instructions, str):
             instructions = [instructions]
 
-        logger.info(f"[prepare_input] image_feature_keys={image_feature_keys}")
         batch_images: list[list[Image.Image]] | None = None
         for key in image_feature_keys:
             imgs = batch[key]
-            if isinstance(imgs, torch.Tensor):
-                logger.info(
-                    f"[prepare_input] key={key} tensor shape={imgs.shape} dtype={imgs.dtype}"
-                )
             if isinstance(imgs, torch.Tensor) and imgs.ndim == 3:
                 imgs = [imgs]
             key_images = [_to_pil(img) for img in imgs]
@@ -111,9 +105,6 @@ def prepare_input(
         for idx, sample_images in enumerate(batch_images):
             batch_images[idx] = [img for img in sample_images if img is not None]
 
-        logger.info(
-            f"[prepare_input] batch_size={len(batch_images)} images_per_sample={[len(s) for s in batch_images]} pil_size={batch_images[0][0].size if batch_images else None}"
-        )
         return batch_images, instructions
 
     def build_qwenvl_inputs(
@@ -139,20 +130,13 @@ def _build_messages(
         return messages
 
     def forward(self, batch: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]:
-        logger.info(
-            f"[VLM.forward] input keys={list(batch.keys())} "
-            + " ".join(f"{k}={v.shape}" for k, v in batch.items() if isinstance(v, torch.Tensor))
-        )
         with torch.autocast(get_platform().amp_device_type(), dtype=torch.bfloat16):
             outputs = self.model(
                 **batch,
                 output_hidden_states=True,
                 return_dict=True,
                 **kwargs,
             )
-        logger.info(
-            f"[VLM.forward] hidden_states: {len(outputs.hidden_states)} layers, last={outputs.hidden_states[-1].shape}"
-        )
         # TODO: (yupu) We should output the original outputs, not just the hidden states.
         return {"hidden_states": outputs.hidden_states}
 
@@ -197,13 +181,6 @@ def build_qwenvl_inputs(
             text=texts, images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt"
         )
 
-        logger.info(
-            "[Qwen25.build_qwenvl_inputs] "
-            + " ".join(
-                f"{k}={v.shape}" for k, v in batch_input.items() if isinstance(v, torch.Tensor)
-            )
-        )
-
         # Use current CUDA device instead of self.model.device, which returns
         # a DTensor device under FSDP2 and causes mixed Tensor/DTensor errors.
         return batch_input.to(get_platform().device())
@@ -245,13 +222,6 @@ def build_qwenvl_inputs(
             return_tensors="pt",
         )
 
-        logger.info(
-            "[Qwen3.build_qwenvl_inputs] "
-            + " ".join(
-                f"{k}={v.shape}" for k, v in batch_inputs.items() if isinstance(v, torch.Tensor)
-            )
-        )
-
         # Use current CUDA device instead of self.model.device, which returns
         # a DTensor device under FSDP2 and causes mixed Tensor/DTensor errors.
         return batch_inputs.to(get_platform().device())