Move date preprocessing into Energon for RoboBrain-X0 training. (flagos-ai#1069)

### PR Category
Train

### PR Types
Improvements

### PR Description
Move date preprocessing into Energon for RoboBrain-X0 training.

Author: Hchnr

examples/robobrain_x0_5/conf/train/libero_qwengroot.yaml

@@ -14,6 +14,9 @@ output_directory: outputs/libero_qwengroot/checkpoints/ckpt_out
 log_freq: 10
 train_steps: 100
 
+shuffle_buffer_size: null
+shuffle_over_epochs_multiplier: -1
+
 framework:
   name: QwenGR00T
   qwenvl:

flagscale/models/robobrain_x/qwen_groot.py

@@ -65,19 +65,7 @@ def __init__(self, config: dict | None = None, **kwargs) -> None:
         self.past_action_window_size = config.framework.action_model.past_action_window_size
         self.chunk_len = self.past_action_window_size + 1 + self.future_action_window_size
 
-    def forward(self, examples: list[dict] | None = None, **kwargs) -> tuple:
-        batch_images = [example["image"] for example in examples]  #  [B，[PLT]]
-        instructions = [example["lang"] for example in examples]  # [B, str]
-        actions = [example["action"] for example in examples]  # label [B， len, 7]
-
-        state = (
-            [example["state"] for example in examples] if "state" in examples[0] else None
-        )  # [B, 1, state_dim]
-
-        # Step 1: QWenVL input format
-        qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(
-            images=batch_images, instructions=instructions
-        )
+    def forward(self, qwen_inputs, state, actions, **kwargs) -> tuple:
         with torch.autocast("cuda", dtype=torch.bfloat16):
             qwenvl_outputs = self.qwen_vl_interface(
                 **qwen_inputs, output_attentions=False, output_hidden_states=True, return_dict=True
@@ -88,14 +76,7 @@ def forward(self, examples: list[dict] | None = None, **kwargs) -> tuple:
         # Step 4: Action Expert Forward and Loss
         with torch.autocast("cuda", dtype=torch.float32):
             # [B, T_full, action_dim]
-            if isinstance(actions[0], torch.Tensor):
-                actions = torch.stack(actions, dim=0).to(
-                    device=last_hidden.device, dtype=last_hidden.dtype
-                )
-            else:
-                actions = torch.tensor(
-                    np.array(actions), device=last_hidden.device, dtype=last_hidden.dtype
-                )
+            actions = actions.to(device=last_hidden.device, dtype=last_hidden.dtype)
             actions_target = actions[
                 :, -(self.future_action_window_size + 1) :, :
             ]  # (B, chunk_len, action_dim)
@@ -109,16 +90,8 @@ def forward(self, examples: list[dict] | None = None, **kwargs) -> tuple:
             last_hidden_repeated = last_hidden.repeat(repeated_diffusion_steps, 1, 1)
 
             state_repeated = None
-            if state is not None:
-                if isinstance(state[0], torch.Tensor):
-                    state = torch.stack(state, dim=0).to(
-                        device=last_hidden.device, dtype=last_hidden.dtype
-                    )
-                else:
-                    state = torch.tensor(
-                        np.array(state), device=last_hidden.device, dtype=last_hidden.dtype
-                    )
-                state_repeated = state.repeat(repeated_diffusion_steps, 1, 1)
+            state = state.to(device=last_hidden.device, dtype=last_hidden.dtype)
+            state_repeated = state.repeat(repeated_diffusion_steps, 1, 1)
 
             action_loss = self.action_model(
                 last_hidden_repeated, actions_target_repeated, state_repeated
@@ -297,38 +270,6 @@ def dryrun_with_random_sample(cfg):
     model.save_pretrained()
 
 
-def dryrun_with_dataloader(cfg):
-    model: Qwen_GR00T = Qwen_GR00T(cfg)
-
-    from megatron.energon import WorkerConfig, get_loader, get_train_dataset
-    from tools.datasets.vla.data.dataset_helpers_np_pil import TaskEncoder
-
-    ds = get_train_dataset(
-        cfg.datasets.data_path,
-        batch_size=1,
-        shuffle_buffer_size=100,
-        max_samples_per_sequence=100,
-        worker_config=WorkerConfig.default_worker_config(num_workers=1, data_parallel_group=None),
-        task_encoder=TaskEncoder(cfg.datasets.task_encoder),
-        repeat=True,
-    )
-    vla_train_dataloader = get_loader(ds)
-    data_iter = iter(vla_train_dataloader)
-    batch = next(data_iter)
-    batch = get_batch(batch)
-
-    # try get model
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device)
-    model(batch)
-    forward_output = model(batch)
-    action_loss = forward_output["action_loss"]
-    print(f"Action Loss: {action_loss.item()}")
-
-    action = model.predict_action(batch_images=[batch[0]["image"]], instructions=[batch[0]["lang"]])
-    print(f"Action inference: {action['normalized_actions'].shape}")
-
-
 if __name__ == "__main__":
     import argparse
 
@@ -339,13 +280,7 @@ def dryrun_with_dataloader(cfg):
         default="./examples/robobrain_x0_5/conf/train/libero_qwengroot.yaml",
         help="Path to YAML config",
     )
-    parser.add_argument("--dryrun-dataloader", action="store_true")
-    parser.add_argument("--dryrun-random", action="store_true")
 
     args, clipargs = parser.parse_known_args()
     cfg = OmegaConf.load(args.config_yaml)
-
-    if args.dryrun_dataloader:
-        dryrun_with_dataloader(cfg)
-    if args.dryrun_random:
-        dryrun_with_random_sample(cfg)
+    dryrun_with_random_sample(cfg)

flagscale/runner/runner_train.py

@@ -464,6 +464,7 @@ def run(
         monitor=False,
         interval=10,
         enable_monitoring=None,
+        **kwargs,
     ):
         # Read from config if not explicitly provided
         if enable_monitoring is None:

flagscale/train/megatron/train_robobrain_x0.5_qwengroot.py

@@ -11,6 +11,7 @@
 import pathlib
 import platform
 import random
+import time
 
 import epath
 import numpy as np
@@ -22,10 +23,10 @@
 
 import wandb
 from megatron.energon import WorkerConfig, get_loader, get_train_dataset
-from tools.datasets.vla.data.dataset_helpers_np_pil import TaskEncoder
+from tools.datasets.vla.data.dataset_helpers_preprocess import TaskEncoder
 
 from flagscale.logger import logger
-from flagscale.models.robobrain_x.qwen_groot import Qwen_GR00T, get_batch
+from flagscale.models.robobrain_x.qwen_groot import Qwen_GR00T
 
 # Sane Defaults
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -94,15 +95,20 @@ def setup_optimizer_and_scheduler(
 
 
 def init_ddp(seed):
-    torch.manual_seed(seed)
-    np.random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
     random.seed(seed)
-    local_rank = int(os.environ["LOCAL_RANK"])
-    torch.cuda.set_device(local_rank)
-    torch.distributed.init_process_group(backend='nccl', init_method='env://')
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
     torch.backends.cudnn.enabled = True
-    torch.backends.cudnn.benchmark = True
+    torch.backends.cudnn.benchmark = False
     torch.backends.cudnn.deterministic = True
+    torch.use_deterministic_algorithms(True)
+
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    torch.cuda.set_device(local_rank)
+    torch.distributed.init_process_group(backend="nccl", init_method="env://")
     return local_rank
 
 
@@ -126,16 +132,19 @@ def init_wandb(config, *, resuming: bool, log_code: bool = False, enabled: bool
 
 
 def main(cfg) -> None:
+    local_rank = init_ddp(cfg.seed)
+
     # build model
     vla = Qwen_GR00T(cfg)
     # prepare data
     ds = get_train_dataset(
         cfg.datasets.data_path,
         batch_size=cfg.batch_size,
-        shuffle_buffer_size=100,
+        shuffle_buffer_size=cfg.shuffle_buffer_size,
         max_samples_per_sequence=100,
+        shuffle_over_epochs_multiplier=cfg.shuffle_over_epochs_multiplier,
         worker_config=WorkerConfig.default_worker_config(num_workers=1, data_parallel_group=None),
-        task_encoder=TaskEncoder(cfg.datasets.task_encoder),
+        task_encoder=TaskEncoder(cfg),
         repeat=True,
     )
     vla_train_dataloader = get_loader(ds)
@@ -145,13 +154,9 @@ def main(cfg) -> None:
     # set optimizer and scheduler
     optimizer, lr_scheduler = setup_optimizer_and_scheduler(model=vla, cfg=cfg)
     # Run VLA Training
-    local_rank = init_ddp(cfg.seed)
+
     if dist.get_rank() == 0 and local_rank == 0:
         logger.info(f"Running on: {platform.node()}")
-        if cfg.batch_size % torch.cuda.device_count() != 0:
-            raise ValueError(
-                f"Batch size {cfg.batch_size} must be divisible by the number of devices {torch.cuda.device_count()}."
-            )
         resuming = cfg.resume
         init_wandb(cfg, resuming=resuming, enabled=cfg.wandb_enabled)
 
@@ -160,18 +165,28 @@ def main(cfg) -> None:
 
     step = 0
     done = False
+
+    t_start = time.time()
     while not done:
         batch = next(data_iter)
-        batch = get_batch(batch)
-        output_dict = vla.forward(batch)
+
+        qwen_inputs, state, actions = batch.get("qwen_inputs"), batch.get("state"), batch.get("actions")
+        if not qwen_inputs or not actions:
+            continue
+        for i in qwen_inputs:
+            qwen_inputs[i] = qwen_inputs[i].to(device=vla.device)
+        output_dict = vla.forward(qwen_inputs=qwen_inputs, state=state, actions=actions)
+
         action_loss = output_dict["action_loss"]
         action_loss.backward()
         optimizer.step()
         lr_scheduler.step()
         optimizer.zero_grad()
 
-        if step % cfg.log_freq == 0:
-            logger.info(f"step: {step} loss: {action_loss.item():.3f}")
+        if step % cfg.log_freq == 0 and dist.get_rank() == 0 and local_rank == 0:
+            logger.info(f"step {step} loss: {action_loss.item()}")
+            logger.info(f"step {step}: {(time.time() - t_start) / cfg.log_freq:.3f}s/iter")
+            t_start = time.time()
         step += 1
         if step >= cfg.train_steps:
             done = True

flagscale/train/train_pi.py

@@ -76,7 +76,7 @@ def set_seed(seed: int):
 
 
 def init_ddp():
-    local_rank = int(os.environ["LOCAL_RANK"])
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
     torch.cuda.set_device(local_rank)
     torch.distributed.init_process_group(backend="nccl", init_method="env://")