[lite] [refactor] Rename the arealite folder to areal

.github/workflows/test-arealite.yml → .github/workflows/test-areal.yml

@@ -1,20 +1,20 @@
-name: Test AReaLite
+name: Test AReaL
 
 on:
   pull_request:
     branches: [ lite ]
     paths:
-      - .github/workflows/test-arealite.yml
-      - arealite/**
+      - .github/workflows/test-areal.yml
+      - areal/**
       - ci/**
   workflow_dispatch:
 
 jobs:
-  test-arealite:
+  test-areal:
     environment:
-      name: AReaLite-unittests
+      name: AReaL-unittests
     runs-on: gpu
     steps:
       - uses: actions/checkout@v4
 
-      - run: bash ./ci/test_arealite.sh
+      - run: bash ./ci/test_areal.sh

.pre-commit-config.yaml

@@ -38,7 +38,7 @@ repos:
           - -r
           - --remove-all-unused-imports
           - --remove-unused-variables
-        files: ^(arealite/|examples/arealite/).*\.py$
+        files: ^(areal/|examples/lite/).*\.py$
 
   - repo: https://github.com/pycqa/isort
     rev: 5.13.2

arealite/README.md → areal/README.md

@@ -1,9 +1,9 @@
-# AReaLite Design Doc
+# AReaL-Lite Design Doc
 
 ## TL;DR
 
-Follow our [step-by-step code walk-through](../docs/arealite/gsm8k_grpo.md) to
-immediately get started with AReaLite!
+Follow our [step-by-step code walk-through](../docs/areal/gsm8k_grpo.md) to immediately
+get started with AReaL-Lite!
 
 ## Motivation
 
@@ -26,13 +26,13 @@ possible to achieve comparable efficiency with significantly fewer lines of code
 presents an ideal opportunity to redesign the API and distill the massive codebase into
 something clean and maintainable. Rather than pursuing maximum efficiency, our goal is
 to deliver 90% of AReaL's functionality while dramatically reducing code complexity and
-user burden. This philosophy drives AReaLite — the lightweight version of AReaL.
+user burden. This philosophy drives AReaL-Lite — the lightweight version of AReaL.
 
-AReaLite serves as the first phase in AReaL's broader refactoring initiative. It
+AReaL-Lite serves as the first phase in AReaL's broader refactoring initiative. It
 functions both as a standalone training library with intuitive interfaces and as the
 foundation for AReaL's future core API definitions. The plan is to transform AReaL's
-current worker-based architecture into an AI-centric architecture similar to AReaLite,
-where AReaL will **extend** AReaLite's APIs and implementations to support additional
+current worker-based architecture into an AI-centric architecture similar to AReaL-Lite,
+where AReaL will **extend** AReaL-Lite's APIs and implementations to support additional
 backends for efficient large-scale training.
 
 ## Design Principles
@@ -55,7 +55,7 @@ Our design is guided by seven core principles:
 ### Core Directory Structure
 
 ```
-arealite/
+areal/
 ├── api/           # Abstract interfaces and dataclasses
 ├── engine/        # Training and inference engines
 ├── launcher/      # Launcher for different backends
@@ -97,17 +97,17 @@ pipelines:
 - **`workflow/rlvr.py`**: RLVR workflow that utilizes an `InferenceEngine` to sample
   multiple responses for each prompt
 
-#### 4. Entry Point Layer (`examples/arealite/`)
+#### 4. Entry Point Layer (`examples/lite/`)
 
 The entry point layer composes customization layer implementations to create complete RL
 training pipelines. While we provide several reference examples, users have complete
 freedom to adapt these to their specific use cases.
 
-Entry points can be launched using launchers from `arealite/launcher/`, similar to other
+Entry points can be launched using launchers from `areal/launcher/`, similar to other
 distributed launch tools like `torchrun`:
 
 ```bash
-python3 -m arealite.launcher.ray entrypoint.py --config my-config.yaml
+python3 -m areal.launcher.ray entrypoint.py --config my-config.yaml
 ```
 
 ## Usage Examples
@@ -119,17 +119,17 @@ parameters for hyperparameter searches or other experimental needs:
 
 ```bash
 # Launch with Ray launcher: 4 nodes (4 GPUs each), 3 nodes for generation, 1 node for training
-python3 -m arealite.launcher.ray examples/arealite/gsm8k_grpo.py \
-    --config examples/arealite/configs/gsm8k_grpo.yaml \
+python3 -m areal.launcher.ray examples/lite/gsm8k_grpo.py \
+    --config examples/lite/configs/gsm8k_grpo.yaml \
     experiment_name=<your_experiment_name> \
     trial_name=<your_trial_name> \
     allocation_mode=sglang.d12p1t1+d4p1t1 \
     cluster.n_nodes=4 \
     cluster.n_gpus_per_node=4
 
 # Launch with Slurm launcher: 16 nodes (8 GPUs each), 12 nodes for generation, 4 nodes for training
-python3 -m arealite.launcher.slurm examples/arealite/gsm8k_grpo.py \
-    --config examples/arealite/configs/gsm8k_grpo.yaml \
+python3 -m areal.launcher.slurm examples/lite/gsm8k_grpo.py \
+    --config examples/lite/configs/gsm8k_grpo.yaml \
     experiment_name=<your_experiment_name> \
     trial_name=<your_trial_name> \
     allocation_mode=sglang.d96p1t1+d32p1t1 \

arealite/api/cli_args.py → areal/api/cli_args.py

@@ -11,7 +11,7 @@
 from hydra import initialize as hydra_init
 from omegaconf import MISSING, OmegaConf
 
-from arealite.utils.fs import get_user_tmp
+from areal.utils.fs import get_user_tmp
 
 
 @dataclass
@@ -189,6 +189,15 @@ class TrainEngineConfig:
     )
     # Runtime microbatch limit
     mb_spec: MicroBatchSpec = field(default_factory=MicroBatchSpec)
+    pad_to_maximum: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to pad each microbatch to the length upper bound specified by mb_spec. "
+                "Can reduce memory fragmentation but slow down training."
+            )
+        },
+    )
 
     # Training Backend Configuration
     disable_dropout: bool = field(default=False)

arealite/api/engine_api.py → areal/api/engine_api.py

@@ -8,7 +8,7 @@
 from tensordict import TensorDict
 from torchdata.stateful_dataloader import StatefulDataLoader
 
-from arealite.api.io_struct import (
+from areal.api.io_struct import (
     FinetuneSpec,
     LLMRequest,
     LLMResponse,
@@ -18,7 +18,7 @@
 )
 
 if TYPE_CHECKING:
-    from arealite.api.workflow_api import RolloutWorkflow
+    from areal.api.workflow_api import RolloutWorkflow
 
 
 @dataclass

arealite/api/io_struct.py → areal/api/io_struct.py

@@ -8,16 +8,14 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple
 
-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
+from areal.api.cli_args import GenerationHyperparameters, SaverConfig
+from areal.utils.network import find_free_ports, gethostip
 
 if TYPE_CHECKING:
-    from arealite.api.engine_api import TrainEngine
+    from areal.api.engine_api import TrainEngine
 
 
 @dataclass
@@ -81,6 +79,7 @@ def steps_per_epoch(self):
 
 
 class AllocationType(enum.Enum):
+    COLOCATE = 0
     DECOUPLED_vLLM = 1
     DECOUPLED_SGLANG = 2
 
@@ -124,11 +123,17 @@ def train_world_size(self) -> int:
 
     @classmethod
     def from_str(cls, allocation_mode: str):
-        alloc_decoupled = AllocationMode.extract_decoupled_alloc(allocation_mode)
         if "vllm" in allocation_mode:
+            alloc_decoupled = AllocationMode.extract_decoupled_alloc(allocation_mode)
             return cls(AllocationType.DECOUPLED_vLLM, alloc_decoupled)
         elif "sglang" in allocation_mode:
+            alloc_decoupled = AllocationMode.extract_decoupled_alloc(allocation_mode)
             return cls(AllocationType.DECOUPLED_SGLANG, alloc_decoupled)
+        else:
+            return cls(
+                AllocationType.COLOCATE,
+                AllocationMode.extract_parallelism_strategy(allocation_mode),
+            )
         raise NotImplementedError(f"Failed to parse allocation: {allocation_mode}")
 
     @staticmethod
@@ -196,7 +201,7 @@ def from_disk(
         cls,
         saver_config: SaverConfig,
     ) -> "WeightUpdateMeta":
-        from arealite.utils.saver import Saver
+        from areal.utils.saver import Saver
 
         path = os.path.join(
             Saver.get_save_checkpoint_root(saver_config),

arealite/api/workflow_api.py → areal/api/workflow_api.py

@@ -11,16 +11,16 @@
 from tensordict import TensorDict
 from torchdata.stateful_dataloader import StatefulDataLoader
 
-from arealite.api.cli_args import InferenceEngineConfig
-from arealite.api.engine_api import InferenceEngine
-from arealite.api.io_struct import RolloutStat
-from arealite.utils.data import concat_padded_tensors
+from areal.api.cli_args import InferenceEngineConfig
+from areal.api.engine_api import InferenceEngine
+from areal.api.io_struct import RolloutStat
+from areal.utils.data import concat_padded_tensors
 from realhf.base import logging
 
 if TYPE_CHECKING:
-    from arealite.api.engine_api import InferenceEngine
+    from areal.api.engine_api import InferenceEngine
 
-logger = logging.getLogger("arealite.workflow_api")
+logger = logging.getLogger("areal.workflow_api")
 
 
 ROLLOUT_POLL_WAIT_TIME = 0.05
@@ -33,7 +33,7 @@ async def arun_episode(
     ) -> TensorDict:
         """Run a single episode of the workflow.
 
-        See concrete example implementations under the `arealite/workflow` directory.
+        See concrete example implementations under the `areal/workflow` directory.
         """
         raise NotImplementedError()
 

arealite/dataset/__init__.py → areal/dataset/__init__.py

@@ -17,25 +17,21 @@ def get_custom_dataset(
 ):
 
     if "gsm8k" in path and type == "sft":
-        from examples.arealite.dataset.gsm8k import get_gsm8k_sft_dataset
+        from areal.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
+        from areal.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,
-        )
+        from areal.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,
-        )
+        from areal.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

examples/arealite/dataset/gsm8k.py → areal/dataset/gsm8k.py

@@ -1,4 +1,4 @@
-from datasets import Dataset, load_dataset
+from datasets import load_dataset
 from datasets.distributed import split_dataset_by_node
 
 

arealite/engine/base_hf_engine.py → areal/engine/base_hf_engine.py

@@ -17,9 +17,9 @@
     get_linear_schedule_with_warmup,
 )
 
-from arealite.api.cli_args import TrainEngineConfig
-from arealite.api.engine_api import FinetuneSpec, TrainEngine
-from arealite.utils.data import (
+from areal.api.cli_args import TrainEngineConfig
+from areal.api.engine_api import FinetuneSpec, TrainEngine
+from areal.utils.data import (
     MicroBatchList,
     amend_position_ids,
     pack_tensor_dict,
@@ -30,8 +30,8 @@
     unpack_sequence,
     unsqueeze_mb_list,
 )
-from arealite.utils.fsdp import get_cosine_schedule_with_warmup
-from arealite.utils.model import VALID_VISION_MODELS, disable_dropout_in_model
+from areal.utils.fsdp import get_cosine_schedule_with_warmup
+from areal.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
 
@@ -259,11 +259,16 @@ def prepare_mb_list(self, input_: TensorDict) -> MicroBatchList:
         input_ = amend_position_ids(input_)
 
         mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec)
+        mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs]
+        mb_list = pad_mb_list(
+            mb_list,
+            pad_value=0.0,
+            pad_to_maximum=self.config.pad_to_maximum,
+        )
         logger.info(
-            f"Microbatch #tokens (rank {dist.get_rank()}): {mb_list.group_lens}"
+            f"Microbatch #tokens (rank {dist.get_rank()}): {mb_list.group_lens}, "
+            f"padded to: {mb_list.padded_to_lengths}, padding lengths: {mb_list.padding_lengths}"
         )
-        mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs]
-        mb_list = pad_mb_list(mb_list, pad_value=0.0)
         # 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)

arealite/engine/fsdp_engine.py → areal/engine/fsdp_engine.py

@@ -13,21 +13,20 @@
 )
 from transformers import AutoProcessor, PreTrainedTokenizerFast
 
-from arealite.api.cli_args import TrainEngineConfig
-from arealite.api.engine_api import FinetuneSpec
-from arealite.api.io_struct import ParamSpec, SaveLoadMeta, WeightUpdateMeta
-from arealite.engine.base_hf_engine import BaseHFEngine
-from arealite.utils.distributed import init_custom_process_group
-from arealite.utils.fsdp import (
+from areal.api.cli_args import TrainEngineConfig
+from areal.api.engine_api import FinetuneSpec
+from areal.api.io_struct import ParamSpec, SaveLoadMeta, WeightUpdateMeta
+from areal.engine.base_hf_engine import BaseHFEngine
+from areal.utils.distributed import init_custom_process_group
+from areal.utils.fsdp import (
     CPUOffloadPolicy,
     MixedPrecisionPolicy,
     apply_fsdp2,
     create_fsdp_device_mesh,
     fsdp2_clip_grad_norm_,
     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 areal.utils.save_load import get_state_dict_from_repo_id_or_path
 from realhf.base import logging, name_resolve, names, pkg_version
 
 logger = logging.getLogger("FSDPEngine")
@@ -46,7 +45,7 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None):
         assert addr is None, "FSDPEngine does not support remote initialization."
         assert pkg_version.is_version_greater_or_equal(
             "torch", "2.4.0"
-        ), f"arealite only supports FSDP2, which requires torch>=2.4.0"
+        ), f"areal only supports FSDP2, which requires torch>=2.4.0"
 
         self.create_process_group()
         self.create_device_model()

arealite/engine/ppo/actor.py → areal/engine/ppo/actor.py

@@ -4,11 +4,11 @@
 import torch
 from tensordict import TensorDict
 
-from arealite.api.cli_args import MicroBatchSpec, PPOActorConfig
-from arealite.api.engine_api import TrainEngine
-from arealite.engine.fsdp_engine import FSDPEngine
-from arealite.utils.data import split_padded_tensor_dict_into_mb_list
-from arealite.utils.functional import (
+from areal.api.cli_args import MicroBatchSpec, PPOActorConfig
+from areal.api.engine_api import TrainEngine
+from areal.engine.fsdp_engine import FSDPEngine
+from areal.utils.data import split_padded_tensor_dict_into_mb_list
+from areal.utils.functional import (
     gather_logprobs,
     gather_logprobs_entropy,
     masked_normalization,

arealite/engine/sft/lm_engine.py → areal/engine/sft/lm_engine.py

@@ -2,10 +2,10 @@
 import torch.utils.data
 from tensordict import TensorDict
 
-from arealite.api.cli_args import TrainEngineConfig
-from arealite.api.engine_api import TrainEngine
-from arealite.engine.fsdp_engine import FSDPEngine
-from arealite.utils.functional import gather_logprobs
+from areal.api.cli_args import TrainEngineConfig
+from areal.api.engine_api import TrainEngine
+from areal.engine.fsdp_engine import FSDPEngine
+from areal.utils.functional import gather_logprobs
 from realhf.base import stats_tracker