feat(fsdp): add memory_efficient_load option for large model initialization

areal/api/cli_args.py

@@ -374,6 +374,16 @@ class FSDPEngineConfig:
         default=False,
         metadata={"help": "Whether to offload FSDP parameters to CPU."},
     )
+    memory_efficient_load: bool = field(
+        default=False,
+        metadata={
+            "help": "Enable memory-efficient model loading. When enabled, model weights "
+            "are initialized on CPU and only rank 0 loads pretrained weights, which are "
+            "then broadcast to all ranks after FSDP sharding. This reduces peak GPU memory "
+            "during initialization for large models. Note: For VLMs, rank 0 broadcast is "
+            "not used; each rank loads weights independently on CPU."
+        },
+    )
 
 
 @dataclass
@@ -840,12 +850,18 @@ class TrainEngineConfig:
     )
 
     def __post_init__(self):
-        """Validate scheduling_spec length."""
+        """Validate scheduling_spec length and config combinations."""
         if len(self.scheduling_spec) not in (1, 2):
             raise ValueError(
                 f"scheduling_spec must contain 1 or 2 SchedulingSpec, "
                 f"got {len(self.scheduling_spec)}"
             )
+        if self.fsdp.memory_efficient_load and self.init_from_scratch:
+            raise ValueError(
+                "memory_efficient_load cannot be used with init_from_scratch=True. "
+                "memory_efficient_load is for loading pretrained weights on CPU, "
+                "but init_from_scratch creates a model without loading any weights."
+            )
 
 
 @dataclass

areal/engine/fsdp_engine.py

@@ -281,12 +281,40 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec, *args, **kwargs):
             CPUOffloadPolicy() if self.config.fsdp.offload_params else None
         )
         tik = time.perf_counter()
-        # Prepare lora weights synchronization
-        if self.config.use_lora:
+
+        full_state = None
+        need_broadcast = False
+
+        is_llm_cpu_load = (
+            self.config.fsdp.memory_efficient_load
+            and not self.config.init_from_scratch
+            and not self.is_vision_model
+        )
+
+        if is_llm_cpu_load or self.config.use_lora:
+            need_broadcast = True
             if dist.get_rank() == 0:
+                if is_llm_cpu_load:
+                    pretrained_state = get_state_dict_from_repo_id_or_path(
+                        self.config.path
+                    )
+                    missing, unexpected = self.model.load_state_dict(
+                        pretrained_state, strict=False
+                    )
+                    if missing:
+                        self.logger.warning(
+                            f"Missing keys when loading pretrained weights: {missing}"
+                        )
+                    if unexpected:
+                        self.logger.warning(
+                            f"Unexpected keys when loading pretrained weights: {unexpected}"
+                        )
+                    del pretrained_state
+                    gc.collect()
                 full_state = self.model.state_dict()
             else:
                 full_state = {}
+
         # NOTE: This applies FSDP2 with N-D parallelism (DP+SP+TP)
         parallelize_model(
             self.model,
@@ -297,14 +325,19 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec, *args, **kwargs):
             cpu_offload=self.cpu_offload,
             wrap_policy=self.config.fsdp.wrap_policy,
         )
-        # Synchronize initialized lora weights
-        if self.config.use_lora:
+
+        if need_broadcast:
+            broadcast_tik = time.perf_counter()
             fsdp2_load_full_state_dict(
                 self.model,
                 full_state,
                 self.cpu_offload,
                 tie_word_embeddings=self.model_config.tie_word_embeddings,
             )
+            self.logger.info(
+                f"Broadcasting model weights took {time.perf_counter() - broadcast_tik:.2f} seconds"
+            )
+
         self.logger.info(
             f"Applying FSDP2 with N-D parallelism for {time.perf_counter() - tik:.2f} seconds"
         )
@@ -703,9 +736,7 @@ def _create_llm_actor_or_critic(self):
         }
         model_kwargs.update(common_kwargs)
 
-        if self.config.init_from_scratch:
-            # initialize model from config
-            # NOTE: VLM cannot directly load state dict using this random initialized model
+        if self.config.init_from_scratch or self.config.fsdp.memory_efficient_load:
             model = model_class.from_config(
                 self.model_config,
                 **model_kwargs,
@@ -724,6 +755,13 @@ def _create_device_model(self):
 
         dtype = getattr(torch, self.config.dtype)
 
+        if self.config.fsdp.memory_efficient_load:
+            loading_device = "cpu"
+        else:
+            loading_device = current_platform.device_type
+
+        self.get_device_stats().log("before model creation/loading")
+
         if self.is_vision_model:
             if dtype == torch.float16:
                 raise ValueError(
@@ -738,8 +776,8 @@ def _create_device_model(self):
             )
 
             tik = time.perf_counter()
-            device = current_platform.device_type
-            with torch.device(device):
+            # VLM: Use from_pretrained() on loading_device.
+            with torch.device(loading_device):
                 model = AutoModelForImageTextToText.from_pretrained(
                     pretrained_model_name_or_path=self.config.path,
                     trust_remote_code=True,
@@ -752,17 +790,20 @@ def _create_device_model(self):
             self.tokenizer = load_hf_tokenizer(self.config.path)
             self.processor = None
             tik = time.perf_counter()
-            with torch.device(current_platform.device_type):
+            # LLM: Decide between from_config() or from_pretrained() based on memory_efficient_load
+            with torch.device(loading_device):
                 model = self._create_llm_actor_or_critic()
                 if self.config.disable_dropout:
                     disable_dropout_in_model(model)
 
+        self.get_device_stats().log("after model creation/loading")
+
         if self.config.gradient_checkpointing:
             model.gradient_checkpointing_enable(
                 gradient_checkpointing_kwargs={"use_reentrant": False}
             )
         self.logger.info(
-            f"Model creation and loading time: {time.perf_counter() - tik}"
+            f"Model creation and loading time: {time.perf_counter() - tik:.2f}s"
         )
         self.model = model
 

areal/tests/test_fsdp_memory_efficient_lora.py

@@ -0,0 +1,41 @@
+import subprocess
+
+import pytest
+
+from areal.api.alloc_mode import AllocationMode
+from areal.platforms import current_platform
+from areal.utils.network import find_free_ports
+
+
+def _run_test_with_torchrun(alloc_mode: str, output: str):
+    port = find_free_ports(1)[0]
+    n_gpus = AllocationMode.from_str(alloc_mode).train.world_size
+    try:
+        subprocess.run(
+            [
+                "torchrun",
+                f"--nproc_per_node={n_gpus}",
+                "--nnodes=1",
+                "--master-addr=localhost",
+                f"--master_port={port}",
+                "areal/tests/torchrun/run_fsdp_memory_efficient_lora.py",
+                f"--allocation_mode={alloc_mode}",
+                f"--output={output}",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except subprocess.CalledProcessError as e:
+        pytest.fail(f"Test failed with error: {e.stderr}, {e.stdout}")
+    with open(output) as f:
+        result = f.read().strip()
+    assert result == "Passed", f"Test failed: {result}"
+
+
+@pytest.mark.slow
+def test_fsdp_memory_efficient_lora(tmp_path_factory):
+    if current_platform.device_count() < 1:
+        pytest.skip("Test requires at least 1 GPU")
+    output = tmp_path_factory.mktemp("test_output") / "fsdp_memory_efficient_lora.out"
+    _run_test_with_torchrun("d1t1c1", str(output))

areal/tests/torchrun/run_fsdp_memory_efficient_lora.py

@@ -0,0 +1,151 @@
+"""Test memory_efficient_load with LoRA configuration."""
+
+import argparse
+import os
+
+import torch
+import torch.distributed as dist
+
+from areal.api.alloc_mode import AllocationMode
+from areal.api.cli_args import (
+    FSDPEngineConfig,
+    MicroBatchSpec,
+    OptimizerConfig,
+    TrainEngineConfig,
+)
+from areal.api.io_struct import FinetuneSpec
+from areal.engine.fsdp_engine import FSDPEngine
+from areal.platforms import current_platform
+from areal.tests.utils import get_model_path
+
+MODEL_PATH = get_model_path(
+    "/storage/openpsi/models/Qwen__Qwen3-0.6B/", "Qwen/Qwen3-0.6B"
+)
+
+
+def write_result(out: str, succ: bool):
+    with open(out, "w") as f:
+        if succ:
+            f.write("Passed")
+        else:
+            f.write("Failed")
+
+
+def make_fsdp_engine_with_lora(
+    allocation_mode: str,
+    memory_efficient_load: bool,
+):
+    """Create FSDPEngine with LoRA and optionally memory_efficient_load."""
+    config = TrainEngineConfig(
+        experiment_name="test_fsdp_memory_efficient_lora",
+        trial_name="test",
+        mb_spec=MicroBatchSpec(max_tokens_per_mb=256),
+        path=MODEL_PATH,
+        optimizer=OptimizerConfig(),
+        fsdp=FSDPEngineConfig(memory_efficient_load=memory_efficient_load),
+        # LoRA config
+        use_lora=True,
+        lora_rank=8,
+        lora_alpha=16,
+        peft_type="lora",
+    )
+    alloc_mode = AllocationMode.from_str(allocation_mode)
+    engine = FSDPEngine(config)
+    ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=128, train_batch_size=8)
+    engine.create_process_group(parallel_strategy=alloc_mode.train)
+    engine.initialize(None, ft_spec)
+    return engine
+
+
+def test_memory_efficient_lora(alloc_mode: str, output: str | None = None):
+    """Test that memory_efficient_load works correctly with LoRA.
+
+    This test verifies:
+    1. Engine initializes successfully with memory_efficient_load=True and use_lora=True
+    2. LoRA layers are properly applied
+    3. Model can perform a forward pass
+    """
+    rank = int(os.environ["RANK"])
+    print(f"Running memory_efficient_load + LoRA test on rank {rank}")
+
+    succ = True
+
+    # Test 1: Create engine with memory_efficient_load=True
+    print(f"Rank {rank}: Creating engine with memory_efficient_load=True")
+    engine = make_fsdp_engine_with_lora(alloc_mode, memory_efficient_load=True)
+
+    # Test 2: Verify LoRA layers exist
+    lora_params = [
+        name for name, _ in engine.model.named_parameters() if "lora" in name.lower()
+    ]
+    if not lora_params:
+        print(f"Rank {rank}: ERROR - No LoRA parameters found!")
+        succ = False
+    else:
+        print(f"Rank {rank}: Found {len(lora_params)} LoRA parameters")
+
+    # Test 3: Verify trainable params are only LoRA params
+    trainable_params = [
+        name for name, p in engine.model.named_parameters() if p.requires_grad
+    ]
+    non_lora_trainable = [p for p in trainable_params if "lora" not in p.lower()]
+    if non_lora_trainable:
+        print(
+            f"Rank {rank}: WARNING - Found non-LoRA trainable params: {non_lora_trainable[:5]}"
+        )
+
+    # Test 4: Simple forward pass to verify model works
+    print(f"Rank {rank}: Testing forward pass")
+    try:
+        with torch.no_grad():
+            engine.eval()
+            # Create simple input
+            input_ids = torch.randint(
+                100, 1000, (2, 32), dtype=torch.long, device=engine.device
+            )
+            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+
+            # Get logits through the model
+            outputs = engine.model(
+                input_ids=input_ids, attention_mask=attention_mask, use_cache=False
+            )
+            logits = outputs.logits
+            print(f"Rank {rank}: Forward pass successful, logits shape: {logits.shape}")
+    except Exception as e:
+        print(f"Rank {rank}: ERROR - Forward pass failed: {e}")
+        succ = False
+
+    current_platform.synchronize()
+    dist.barrier()
+
+    engine.destroy()
+
+    if rank == 0 and output:
+        write_result(output, succ)
+
+    print(
+        f"Rank {rank}: memory_efficient_load + LoRA test {'PASSED' if succ else 'FAILED'}"
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Test memory_efficient_load with LoRA")
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="Optional path to save the output result",
+    )
+    parser.add_argument(
+        "--allocation_mode",
+        type=str,
+        default="d1t1c1",
+        help="Allocation mode for the model",
+    )
+    args = parser.parse_args()
+
+    test_memory_efficient_lora(alloc_mode=args.allocation_mode, output=args.output)
+
+
+if __name__ == "__main__":
+    main()

docs/best_practices/handling_oom.md

@@ -135,6 +135,33 @@ When encountering an OOM error, you can switch to a more memory-efficient optimi
 setting `actor.optimizer.type: <name>` in your YAML configuration file (e.g.,
 `actor.optimizer.type: sgd`).
 
+### 4. Use Memory-Efficient Model Loading
+
+If you're hitting OOM during model initialization (before training even starts), try
+enabling memory-efficient loading:
+
+```yaml
+actor:
+  fsdp:
+    memory_efficient_load: true
+```
+
+This is especially useful for very large models where loading the full model weights
+directly onto each GPU would exceed memory. When `memory_efficient_load: true` is set:
+
+1. All ranks create model structure on CPU (without loading weights for LLM)
+1. FSDP parallelization is applied
+1. Rank 0 loads pretrained weights and broadcasts to all ranks
+1. Weights are transferred to GPU
+
+This approach trades some initialization time for significantly lower peak GPU memory
+during model loading.
+
+**Note for Vision-Language Models (VLMs):** VLMs do not use the rank 0 broadcast
+optimization. When `memory_efficient_load: true` is set for VLMs, weights are loaded on
+CPU instead of GPU, but each rank loads weights independently. This still reduces GPU
+memory usage during initialization but does not reduce CPU memory or disk/network I/O.
+
 ## Resolving Weight Update OOM Errors
 
 Weight updates can eat up a lot of memory, especially when using NCCL synchronization