From 4047986216ffaeba8a7994d19e6338774f7fb018 Mon Sep 17 00:00:00 2001 From: Wentai Zhang Date: Tue, 27 Jan 2026 17:16:01 +0800 Subject: [PATCH] feat(fsdp): add memory_efficient_load option for large model initialization When training large models, loading weights directly onto each GPU can cause OOM during initialization. This adds `fsdp.memory_efficient_load` config option that initializes model structure on CPU, then rank 0 loads pretrained weights and broadcasts to all ranks after FSDP sharding. This reduces peak GPU memory during model loading at the cost of slightly longer initialization time. --- areal/api/cli_args.py | 18 ++- areal/engine/fsdp_engine.py | 63 ++++++-- .../tests/test_fsdp_memory_efficient_lora.py | 41 +++++ .../run_fsdp_memory_efficient_lora.py | 151 ++++++++++++++++++ docs/best_practices/handling_oom.md | 27 ++++ docs/cli_reference.md | 9 +- 6 files changed, 293 insertions(+), 16 deletions(-) create mode 100644 areal/tests/test_fsdp_memory_efficient_lora.py create mode 100644 areal/tests/torchrun/run_fsdp_memory_efficient_lora.py diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index fe1fd2bcde..7f37e610ec 100644 --- a/areal/api/cli_args.py +++ b/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 diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 330636a3eb..b03e353d10 100644 --- a/areal/engine/fsdp_engine.py +++ b/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 diff --git a/areal/tests/test_fsdp_memory_efficient_lora.py b/areal/tests/test_fsdp_memory_efficient_lora.py new file mode 100644 index 0000000000..63b6774d38 --- /dev/null +++ b/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)) diff --git a/areal/tests/torchrun/run_fsdp_memory_efficient_lora.py b/areal/tests/torchrun/run_fsdp_memory_efficient_lora.py new file mode 100644 index 0000000000..0359a78060 --- /dev/null +++ b/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() diff --git a/docs/best_practices/handling_oom.md b/docs/best_practices/handling_oom.md index 94a71f9654..f881aa7788 100644 --- a/docs/best_practices/handling_oom.md +++ b/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: ` 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 diff --git a/docs/cli_reference.md b/docs/cli_reference.md index 4e0be0e248..ce30433a22 100644 --- a/docs/cli_reference.md +++ b/docs/cli_reference.md @@ -250,10 +250,11 @@ Configuration for Supervised Fine-Tuning (SFT) experiments. Configuration for Fully Sharded Data Parallel (FSDP) training backend. -| Parameter | Type | Default | Description | -| ---------------- | ---------------------------------------------------- | ------- | -------------------------------------------------- | -| `wrap_policy` | [`FSDPWrapPolicy`](section-fsdp-wrap-policy) \| None | `None` | FSDP wrap policy, specifying model layers to wrap. | -| `offload_params` | boolean | `False` | Whether to offload FSDP parameters to CPU. | +| Parameter | Type | Default | Description | +| ----------------------- | ---------------------------------------------------- | ------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `wrap_policy` | [`FSDPWrapPolicy`](section-fsdp-wrap-policy) \| None | `None` | FSDP wrap policy, specifying model layers to wrap. | +| `offload_params` | boolean | `False` | Whether to offload FSDP parameters to CPU. | +| `memory_efficient_load` | boolean | `False` | 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. | (section-fsdp-wrap-policy)=