[WIP][feat] Initial support for VLMs, add Qwen2VL SFT test and Qwen2.5VL GRPO test

arealite/api/cli_args.py

@@ -611,8 +611,15 @@ class ClusterSpecConfig:
 
 @dataclass
 class DatasetConfig:
+    path: str = field(
+        default=MISSING,
+        metadata={
+            "help": "Path to the dataset. Can be a local path or a HuggingFace dataset name."
+        },
+    )
     type: Optional[str] = field(
-        default=None, metadata={"help": "Type of implemented dataset"}
+        default=None,
+        metadata={"help": "Type of training method.e.g., 'sft', 'rl', etc."},
     )
     batch_size: int = field(
         default=1, metadata={"help": "Batch size of the dataloader"}
@@ -743,7 +750,7 @@ class BaseExperimentConfig:
     tokenizer_path: str = field(default="")
 
     train_dataset: DatasetConfig = field(default_factory=DatasetConfig)
-    valid_dataset: DatasetConfig = field(default_factory=DatasetConfig)
+    valid_dataset: Optional[DatasetConfig] = field(default=None)
 
     saver: SaverConfig = field(default_factory=SaverConfig)
     checkpointer: SaverConfig = field(default_factory=SaverConfig)

arealite/api/io_struct.py

@@ -8,7 +8,10 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple
 
-from transformers import PreTrainedTokenizerFast
+import torch
+from gymnasium.core import ActType, ObsType
+from PIL.Image import Image as ImageObject
+from transformers import AutoProcessor, PreTrainedTokenizerFast
 
 from arealite.api.cli_args import GenerationHyperparameters, SaverConfig
 from arealite.utils.network import find_free_ports, gethostip
@@ -51,6 +54,16 @@ def output_len(self) -> int:
         return len(self.output_tokens)
 
 
+@dataclass
+class VLMRequest(LLMRequest):
+    image_data: Optional[List[ImageObject | str]] = field(default_factory=list)
+
+
+@dataclass
+class VLMResponse(LLMResponse):
+    input_images: List[ImageObject | str] = field(default_factory=list)
+
+
 @dataclass
 class FinetuneSpec:
     total_train_epochs: int
@@ -216,7 +229,8 @@ class SaveLoadMeta:
     path: str
     weight_format: str
     with_optim: bool
-    tokenizer: Optional[PreTrainedTokenizerFast]
+    tokenizer: PreTrainedTokenizerFast | None
+    processor: AutoProcessor | None
     base_model_path: str | None
     naive_distributed: bool = False
 

arealite/dataset/__init__.py

@@ -0,0 +1,47 @@
+from typing import Optional
+
+import transformers
+
+VALID_DATASETS = ["gsm8k", "clevr_count_70k"]
+
+
+def get_custom_dataset(
+    path: str,
+    rank: int,
+    world_size: int,
+    type: str = "sft",
+    split: Optional[str] = None,
+    tokenizer: Optional[transformers.PreTrainedTokenizerFast] = None,
+    processor: Optional[transformers.AutoProcessor] = None,
+    **kwargs,
+):
+
+    if "gsm8k" in path and type == "sft":
+        from examples.arealite.dataset.gsm8k import get_gsm8k_sft_dataset
+
+        return get_gsm8k_sft_dataset(path, split, tokenizer, rank, world_size, **kwargs)
+    elif "gsm8k" in path and type == "rl":
+        from examples.arealite.dataset.gsm8k import get_gsm8k_rl_dataset
+
+        return get_gsm8k_rl_dataset(path, split, rank, world_size, **kwargs)
+    elif "clevr_count_70k" in path and type == "sft":
+        from examples.arealite.dataset.clevr_count_70k import (
+            get_clevr_count_70k_sft_dataset,
+        )
+
+        return get_clevr_count_70k_sft_dataset(
+            path, split, processor, rank, world_size, **kwargs
+        )
+    elif "clevr_count_70k" in path and type == "rl":
+        from examples.arealite.dataset.clevr_count_70k import (
+            get_clevr_count_70k_rl_dataset,
+        )
+
+        return get_clevr_count_70k_rl_dataset(
+            path, split, processor, rank, world_size, **kwargs
+        )
+    else:
+        raise ValueError(
+            f"Dataset {path} with split {split} and training type {type} is not supported. "
+            f"Supported datasets are: {VALID_DATASETS}. "
+        )

arealite/engine/base_hf_engine.py

@@ -9,6 +9,8 @@
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoProcessor,
     PretrainedConfig,
     PreTrainedTokenizerFast,
     get_constant_schedule_with_warmup,
@@ -29,8 +31,8 @@
     unsqueeze_mb_list,
 )
 from arealite.utils.fsdp import get_cosine_schedule_with_warmup
-from arealite.utils.model import disable_dropout_in_model
-from realhf.api.core.data_api import load_hf_tokenizer
+from arealite.utils.model import VALID_VISION_MODELS, disable_dropout_in_model
+from realhf.api.core.data_api import load_hf_processor_and_tokenizer, load_hf_tokenizer
 from realhf.base import constants, logging
 
 logger = logging.getLogger("Base HF Engine")
@@ -44,6 +46,7 @@ def __init__(self, config: TrainEngineConfig):
         self.model: torch.nn.Module
         self.optimizer: torch.optim.Optimizer
         self.tokenizer: PreTrainedTokenizerFast
+        self.processor: AutoProcessor | None = None
         # huggingface model config
         self.model_config: PretrainedConfig
         self._version: int = 0
@@ -54,6 +57,12 @@ def __init__(self, config: TrainEngineConfig):
         self._parallelism_group: dist.ProcessGroup
         self.weight_update_group_initialized = False
 
+        self.model_config = AutoConfig.from_pretrained(
+            pretrained_model_name_or_path=self.config.path,
+            trust_remote_code=True,
+        )
+        self.is_vision_model = self.model_config.model_type in VALID_VISION_MODELS
+
         self.world_size = int(os.environ["WORLD_SIZE"])
 
     def set_version(self, version: int):
@@ -92,32 +101,54 @@ def create_device_model(self):
         self.device = torch.device(int(os.environ["LOCAL_RANK"]))
 
         dtype = getattr(torch, self.config.dtype)
-        self.model_config = AutoConfig.from_pretrained(
-            pretrained_model_name_or_path=self.config.path,
-            trust_remote_code=True,
-        )
-        self.tokenizer = load_hf_tokenizer(self.config.path)
-        tik = time.perf_counter()
-        with torch.device("cuda"):
+
+        if self.is_vision_model:
+            if dtype == torch.float16:
+                raise ValueError(
+                    "Vision models do not support float16 dtype. Please use bfloat16."
+                )
             if self.config.init_from_scratch:
-                # initialize scratch model from config
-                # NOTE: VLM cannot directly load state dict using this
-                # random initialized model, so otherwise we call
-                # from_pretrained rather than loading weights into this random model.
-                model = AutoModelForCausalLM.from_config(
-                    self.model_config,
-                    torch_dtype=dtype,
-                    attn_implementation=self.config.attn_impl,
+                raise ValueError(
+                    "Vision models do not support initialization from scratch. Please use a pretrained model."
                 )
-            else:
-                model = AutoModelForCausalLM.from_pretrained(
+            self.processor, self.tokenizer = load_hf_processor_and_tokenizer(
+                self.config.path
+            )
+
+            tik = time.perf_counter()
+            with torch.device("cuda"):
+                model = AutoModelForImageTextToText.from_pretrained(
                     pretrained_model_name_or_path=self.config.path,
                     trust_remote_code=True,
                     torch_dtype=dtype,
                     attn_implementation=self.config.attn_impl,
                 )
-            if self.config.disable_dropout:
-                disable_dropout_in_model(model)
+                if self.config.disable_dropout:
+                    disable_dropout_in_model(model)
+        else:
+            self.tokenizer = load_hf_tokenizer(self.config.path)
+            tik = time.perf_counter()
+            with torch.device("cuda"):
+                if self.config.init_from_scratch:
+                    # initialize scratch model from config
+                    # NOTE: VLM cannot directly load state dict using this
+                    # random initialized model, so otherwise we call
+                    # from_pretrained rather than loading weights into this random model.
+                    model = AutoModelForCausalLM.from_config(
+                        self.model_config,
+                        torch_dtype=dtype,
+                        attn_implementation=self.config.attn_impl,
+                    )
+                else:
+                    model = AutoModelForCausalLM.from_pretrained(
+                        pretrained_model_name_or_path=self.config.path,
+                        trust_remote_code=True,
+                        torch_dtype=dtype,
+                        attn_implementation=self.config.attn_impl,
+                    )
+                if self.config.disable_dropout:
+                    disable_dropout_in_model(model)
+
         if self.config.gradient_checkpointing:
             model.gradient_checkpointing_enable(
                 gradient_checkpointing_kwargs={"use_reentrant": False}
@@ -218,9 +249,15 @@ def step_lr_scheduler(self):
 
     def prepare_mb_list(self, input_: TensorDict) -> MicroBatchList:
         assert "attention_mask" in input_ and "input_ids" in input_
+        if self.is_vision_model:
+            assert (
+                "pixel_values" in input_ and "image_grid_thw" in input_
+            ), "For vision-language models, pixel_values and image_grid_thw must be present in input_"
+
         if isinstance(input_, dict):
             input_ = TensorDict(input_, batch_size=[input_["input_ids"].shape[0]])
         input_ = amend_position_ids(input_)
+
         mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec)
         logger.info(
             f"Microbatch #tokens (rank {dist.get_rank()}): {mb_list.group_lens}"
@@ -230,13 +267,15 @@ def prepare_mb_list(self, input_: TensorDict) -> MicroBatchList:
         # NOTE: We unsqueeze here because huggingface transformer models requires
         # packed input to be of shape [1, total_seqlen].
         mb_list = unsqueeze_mb_list(mb_list)
+
         # FIXME: the resulting max_seqlen is a tensor rather than an integer
         for mb in mb_list.mbs:
             mb["max_seqlen"] = int(mb["max_seqlen"])
             mb["use_cache"] = False
         for mb in mb_list.padded_mbs:
             mb["max_seqlen"] = int(mb["max_seqlen"])
             mb["use_cache"] = False
+
         return mb_list
 
     def train_batch(
@@ -264,11 +303,13 @@ def train_batch(
         for i, (pad_length, padded_mb_input, mb_input) in enumerate(
             zip(mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs)
         ):
+
             outputs = self.model(**padded_mb_input)
 
             logits = outputs.logits.squeeze(0)
             logits = logits[:-pad_length] if pad_length > 0 else logits
             loss = loss_fn(logits, mb_input)
+
             loss_scale = loss_weight_fn(mb_input) / total_loss_weight
 
             # Scale loss for accumulation

arealite/engine/fsdp_engine.py

@@ -11,7 +11,7 @@
     StateDictOptions,
     get_model_state_dict,
 )
-from transformers import PreTrainedTokenizerFast
+from transformers import AutoProcessor, PreTrainedTokenizerFast
 
 from arealite.api.cli_args import TrainEngineConfig
 from arealite.api.engine_api import FinetuneSpec
@@ -27,6 +27,7 @@
     fsdp2_load_full_state_dict,
 )
 from arealite.utils.save_load import get_state_dict_from_repo_id_or_path
+from realhf.api.core.data_api import load_hf_processor_and_tokenizer
 from realhf.base import logging, name_resolve, names, pkg_version
 
 logger = logging.getLogger("FSDPEngine")
@@ -77,7 +78,7 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None):
 
     def save(self, meta: SaveLoadMeta):
         if meta.weight_format == "hf":
-            self._save_model_to_hf(meta.path, meta.tokenizer)
+            self._save_model_to_hf(meta.path, meta.tokenizer, meta.processor)
         elif meta.weight_format == "dcp":
             # TODO: implement DCP save/load for FSDP
             raise NotImplementedError("DCP format saving is not implemented yet. ")
@@ -100,7 +101,10 @@ def load(self, meta: SaveLoadMeta):
             self.load_optimizer_state(meta.path)
 
     def _save_model_to_hf(
-        self, path: str, tokenizer: Optional[PreTrainedTokenizerFast]
+        self,
+        path: str,
+        tokenizer: Optional[PreTrainedTokenizerFast],
+        processor: Optional[AutoProcessor],
     ):
         """Save model in HuggingFace format."""
         if self.model is None:
@@ -119,6 +123,8 @@ def _save_model_to_hf(
             self.model_config.save_pretrained(path)
             if tokenizer is not None:
                 tokenizer.save_pretrained(path)
+            if processor is not None:
+                processor.save_pretrained(path)
 
         dist.barrier(device_ids=[self.device.index])
 
@@ -144,13 +150,13 @@ def upload_weights(self, meta: WeightUpdateMeta):
             dist.barrier(device_ids=[self.device.index])
             torch.cuda.synchronize()
         elif meta.type == "disk":
-            self._save_model_to_hf(meta.path, self.tokenizer)
+            self._save_model_to_hf(meta.path, self.tokenizer, self.processor)
             # dist.barrier() are called when _save_model_to_hf finished
             if dist.get_rank() == 0:
                 update_name = names.update_weights_from_disk(
                     self.config.experiment_name,
                     self.config.trial_name,
-                    self.model_version,
+                    self.get_version(),
                 )
                 name_resolve.add(
                     update_name, str(datetime.now().timestamp()), keepalive_ttl=120
@@ -247,9 +253,11 @@ def train_batch(
             loss.backward()
 
         # NOTE: grad norm clip function is different
+
         grad_norm = fsdp2_clip_grad_norm_(
             self.model.parameters(), max_norm=self.optimizer_config.gradient_clipping
         )
+
         if not torch.isfinite(grad_norm):
             self.optimizer.zero_grad()
             update_successful = False

arealite/engine/sft/lm_engine.py

@@ -52,7 +52,6 @@ def compute_packed_sft_loss(logits: torch.Tensor, input_: TensorDict) -> torch.T
     logprobs = torch.where(loss_mask, logprobs, 0)
 
     loss = -logprobs.sum() / loss_mask.count_nonzero()
-
     with torch.no_grad():
         seqlogp = torch.zeros(
             cu_seqlens.shape[0] - 1, device=logits.device, dtype=torch.float64