From 030c79d652cd43fe4d79f411619677173eb82a64 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 13:04:27 -0700 Subject: [PATCH 1/9] refactor(fsdp): [mechanical] extract actor model loading --- miles/backends/fsdp_utils/actor.py | 192 +---------------- .../backends/fsdp_utils/input_dtype_policy.py | 2 +- miles/backends/fsdp_utils/model_loader.py | 200 ++++++++++++++++++ .../_param_dtype_map_integration_worker.py | 2 +- 4 files changed, 208 insertions(+), 188 deletions(-) create mode 100644 miles/backends/fsdp_utils/model_loader.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f3eab3694..1ec425904 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -1,7 +1,6 @@ import logging -import warnings from argparse import Namespace -from contextlib import contextmanager, nullcontext +from contextlib import nullcontext import ray import torch @@ -39,9 +38,9 @@ from .loss_hub import DiffusionLossContext, flow_grpo_loss_formula, prepare_flow_grpo_batch from .lr_scheduler import get_lr_scheduler from .metrics import new_metric_buffer -from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str +from .mixed_precision import parse_dtype_from_str +from .model_loader import load_fsdp_models from .parallel import create_fsdp_parallel_state -from .sequence_parallel.plan import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -106,62 +105,9 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty # flash-attn is opaque to torch's determinism flag; backends patch their own dispatch. self.model_backend.enable_deterministic_attention(args.fsdp_attention_backend) self.scheduler = self.model_backend.load_scheduler(args) - rank = dist.get_rank() - materialize_weights = rank == 0 - - self.models: dict[str, torch.nn.Module] = {} - for component in args.update_weight_target_modules: - # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap - with self._model_init_context(materialize_weights=materialize_weights): - model = self.model_backend.load_component( - component, - args, - master_dtype=self._master_dtype, - materialize_weights=materialize_weights, - ) - if args.fsdp_attention_backend is not None: - self.model_backend.set_attention_backend(model, args.fsdp_attention_backend) - - # Enable checkpointing on the raw model before PEFT wraps it. The flag - # is consumed when transformer blocks run, so LoRA layers inserted - # below remain inside the checkpointed block forward. - if args.gradient_checkpointing: - self.model_backend.enable_gradient_checkpointing(model) - - if args.use_lora: - model = apply_lora(model, args, self.train_pipeline_config) - - model.train() - - if rank != 0 and any(not parameter.is_meta for parameter in model.parameters()): - raise RuntimeError(f"{component} did not honor meta initialization") - checkpoint.sync_model_dtypes(model) - full_state = model.state_dict() if rank == 0 else {} - model = apply_fsdp2( - model, - self.model_backend.fsdp_parallel_plan(model), - mesh=self.parallel_state.get_mesh("fsdp"), - cpu_offload=self.args.fsdp_cpu_offload, - args=self.args, - ) - checkpoint.broadcast_full_state_to_fsdp( - model, - full_state, - cpu_offload=self.args.fsdp_cpu_offload, - ) - del full_state - self.train_pipeline_config.postprocess_model_after_materialize(model) - self.models[component] = model - - if self.parallel_state.get_optional_mesh("sp") is not None: - for model in self.models.values(): - plan = self.model_backend.sequence_parallel_plan(model) - apply_sequence_parallel( - model, - self.parallel_state, - plan, - self.model_backend.install_sequence_parallel_attention, - ) + self.models = load_fsdp_models( + args, self.model_backend, self.train_pipeline_config, self.parallel_state, master_dtype=self._master_dtype + ) # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() @@ -239,24 +185,6 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty return self.args.start_rollout_id - @contextmanager - def _model_init_context(self, *, materialize_weights: bool): - """Build real CPU weights on rank 0 and meta weights elsewhere.""" - if materialize_weights: - with torch.device("cpu"): - yield - return - - from accelerate import init_empty_weights - - # Some models compute buffer values during __init__, which cannot run on meta. - with init_empty_weights(include_buffers=False), warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=r"for .*: copying from a non-meta parameter in the checkpoint to a meta parameter.*", - ) - yield - @timer def sleep(self) -> None: if not self.args.offload_train: @@ -598,111 +526,3 @@ def move_torch_optimizer(optimizer, device): state[key] = value.to(device, non_blocking=True) torch.cuda.synchronize() - - -def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -> torch.nn.Module: - """Apply PEFT LoRA, leaving non-rank0 adapters uninitialized on meta.""" - from peft import LoraConfig, get_peft_model - - on_meta = dist.get_rank() != 0 - # Per-model fallback when --lora-target-modules is unset (runtime inference: depends on loaded pipeline). - targets = args.lora_target_modules or train_pipeline_config.lora_target_modules - init_lora_weight = args.lora_init_weights - if init_lora_weight == "kaiming-uniform": - init_lora_weight = True # namely kaiming-uniform - model = get_peft_model( - model, - LoraConfig( - r=args.lora_rank, - lora_alpha=args.lora_alpha, - target_modules=targets, - init_lora_weights=False if on_meta else init_lora_weight, - ), - low_cpu_mem_usage=on_meta, - ) - if dist.get_rank() == 0: - model.print_trainable_parameters() - return model - - -def apply_fsdp2( - model, - parallel_plan, - mesh=None, - cpu_offload=False, - args=None, -): - """Apply FSDP2 per the model's FSDPParallelPlan. - - ``parallel_plan.param_dtype_patterns`` is matched against FQNs from ``model``. Each child - ``fully_shard`` call receives exact FQNs relative to that child module, while - parameters managed by the root call retain their root-relative FQNs. - """ - from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard - - offload_policy = CPUOffloadPolicy() if cpu_offload else None - - layer_cls_to_wrap = parallel_plan.no_split_modules - assert layer_cls_to_wrap is not None and len(layer_cls_to_wrap) > 0 and layer_cls_to_wrap[0] is not None - - modules = [module for name, module in model.named_modules() if module.__class__.__name__ in layer_cls_to_wrap] - - param_dtype = parse_dtype_from_str(args.diffusion_forward_dtype) - reduce_dtype = parse_dtype_from_str(args.fsdp_reduce_dtype) - # A wrap entry may also be a module LIST — fully_shard can group several modules into one wrap - # (one shared all-gather); today every wrap holds a single block. - param_dtype_maps = compile_param_dtype_maps( - model, - modules, - parallel_plan.param_dtype_patterns, - param_dtype, - ) - has_param_dtype_overrides = bool(any(param_dtype_maps.wrap_maps) or param_dtype_maps.root_map) - param_dtype_policy_cls = None - if has_param_dtype_overrides: - from .monkey_patches.fsdp_param_dtype_patch import ParamDtypeMixedPrecisionPolicy, apply_param_dtype_map_patch - - apply_param_dtype_map_patch() - param_dtype_policy_cls = ParamDtypeMixedPrecisionPolicy - logger.info( - f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, " - f"param_dtype={param_dtype}, reduce_dtype={reduce_dtype}, " - f"param_dtype_overrides={param_dtype_maps.override_count} " - f"({param_dtype_maps.override_numel:,} parameters)" - ) - - fsdp_kwargs = { - "offload_policy": offload_policy, - "mesh": mesh, - } - - # input_dtype_policy owns boundary casts; autocast owns compute and keeps grad-ckpt recompute consistent. - def make_mp_policy(param_dtype_map): - if param_dtype_map: - assert param_dtype_policy_cls is not None - return param_dtype_policy_cls( - param_dtype=param_dtype, - reduce_dtype=reduce_dtype, - cast_forward_inputs=False, - param_dtype_map=param_dtype_map, - ) - return MixedPrecisionPolicy( - param_dtype=param_dtype, - reduce_dtype=reduce_dtype, - cast_forward_inputs=False, - ) - - for module, wrap_map in zip(modules, param_dtype_maps.wrap_maps, strict=True): - fully_shard( - module, - mp_policy=make_mp_policy(wrap_map), - **fsdp_kwargs, - ) - - fully_shard( - model, - mp_policy=make_mp_policy(param_dtype_maps.root_map), - **fsdp_kwargs, - ) - - return model diff --git a/miles/backends/fsdp_utils/input_dtype_policy.py b/miles/backends/fsdp_utils/input_dtype_policy.py index 966d124f7..9be67a47e 100644 --- a/miles/backends/fsdp_utils/input_dtype_policy.py +++ b/miles/backends/fsdp_utils/input_dtype_policy.py @@ -4,7 +4,7 @@ ("fp32"/"bf16"/"fp16"), "default" for the run's forward dtype, or None to pass the rollout dtype through. The boundary dtype is what element-wise ops see before any weight is involved, so it must match what the family's sglang-d pipeline feeds the DiT for log-prob alignment; compute inside the -model is owned by the trainer's autocast (see actor.apply_fsdp2). +model is owned by the trainer's autocast (see model_loader.apply_fsdp2). """ from __future__ import annotations diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py new file mode 100644 index 000000000..89cf4d019 --- /dev/null +++ b/miles/backends/fsdp_utils/model_loader.py @@ -0,0 +1,200 @@ +import logging +import warnings +from argparse import Namespace +from contextlib import contextmanager + +import torch +import torch.distributed as dist + +from . import checkpoint +from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str +from .sequence_parallel.plan import apply_sequence_parallel + +logger = logging.getLogger(__name__) + + +def load_fsdp_models(args: Namespace, model_backend, train_pipeline_config, parallel_state, *, master_dtype): + rank = dist.get_rank() + materialize_weights = rank == 0 + + models: dict[str, torch.nn.Module] = {} + for component in args.update_weight_target_modules: + # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap + with model_init_context(materialize_weights=materialize_weights): + model = model_backend.load_component( + component, + args, + master_dtype=master_dtype, + materialize_weights=materialize_weights, + ) + if args.fsdp_attention_backend is not None: + model_backend.set_attention_backend(model, args.fsdp_attention_backend) + + # Enable checkpointing on the raw model before PEFT wraps it. The flag + # is consumed when transformer blocks run, so LoRA layers inserted + # below remain inside the checkpointed block forward. + if args.gradient_checkpointing: + model_backend.enable_gradient_checkpointing(model) + + if args.use_lora: + model = apply_lora(model, args, train_pipeline_config) + + model.train() + + if rank != 0 and any(not parameter.is_meta for parameter in model.parameters()): + raise RuntimeError(f"{component} did not honor meta initialization") + checkpoint.sync_model_dtypes(model) + full_state = model.state_dict() if rank == 0 else {} + model = apply_fsdp2( + model, + model_backend.fsdp_parallel_plan(model), + mesh=parallel_state.get_mesh("fsdp"), + cpu_offload=args.fsdp_cpu_offload, + args=args, + ) + checkpoint.broadcast_full_state_to_fsdp( + model, + full_state, + cpu_offload=args.fsdp_cpu_offload, + ) + del full_state + train_pipeline_config.postprocess_model_after_materialize(model) + models[component] = model + + if parallel_state.get_optional_mesh("sp") is not None: + for model in models.values(): + plan = model_backend.sequence_parallel_plan(model) + apply_sequence_parallel( + model, + parallel_state, + plan, + model_backend.install_sequence_parallel_attention, + ) + return models + + +@contextmanager +def model_init_context(*, materialize_weights: bool): + """Build real CPU weights on rank 0 and meta weights elsewhere.""" + if materialize_weights: + with torch.device("cpu"): + yield + return + + from accelerate import init_empty_weights + + # Some models compute buffer values during __init__, which cannot run on meta. + with init_empty_weights(include_buffers=False), warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"for .*: copying from a non-meta parameter in the checkpoint to a meta parameter.*", + ) + yield + + +def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -> torch.nn.Module: + """Apply PEFT LoRA, leaving non-rank0 adapters uninitialized on meta.""" + from peft import LoraConfig, get_peft_model + + on_meta = dist.get_rank() != 0 + # Per-model fallback when --lora-target-modules is unset (runtime inference: depends on loaded pipeline). + targets = args.lora_target_modules or train_pipeline_config.lora_target_modules + init_lora_weight = args.lora_init_weights + if init_lora_weight == "kaiming-uniform": + init_lora_weight = True # namely kaiming-uniform + model = get_peft_model( + model, + LoraConfig( + r=args.lora_rank, + lora_alpha=args.lora_alpha, + target_modules=targets, + init_lora_weights=False if on_meta else init_lora_weight, + ), + low_cpu_mem_usage=on_meta, + ) + if dist.get_rank() == 0: + model.print_trainable_parameters() + return model + + +def apply_fsdp2( + model, + parallel_plan, + mesh=None, + cpu_offload=False, + args=None, +): + """Apply FSDP2 per the model's FSDPParallelPlan. + + ``parallel_plan.param_dtype_patterns`` is matched against FQNs from ``model``. Each child + ``fully_shard`` call receives exact FQNs relative to that child module, while + parameters managed by the root call retain their root-relative FQNs. + """ + from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard + + offload_policy = CPUOffloadPolicy() if cpu_offload else None + + layer_cls_to_wrap = parallel_plan.no_split_modules + assert layer_cls_to_wrap is not None and len(layer_cls_to_wrap) > 0 and layer_cls_to_wrap[0] is not None + + modules = [module for name, module in model.named_modules() if module.__class__.__name__ in layer_cls_to_wrap] + + param_dtype = parse_dtype_from_str(args.diffusion_forward_dtype) + reduce_dtype = parse_dtype_from_str(args.fsdp_reduce_dtype) + # A wrap entry may also be a module LIST — fully_shard can group several modules into one wrap + # (one shared all-gather); today every wrap holds a single block. + param_dtype_maps = compile_param_dtype_maps( + model, + modules, + parallel_plan.param_dtype_patterns, + param_dtype, + ) + has_param_dtype_overrides = bool(any(param_dtype_maps.wrap_maps) or param_dtype_maps.root_map) + param_dtype_policy_cls = None + if has_param_dtype_overrides: + from .monkey_patches.fsdp_param_dtype_patch import ParamDtypeMixedPrecisionPolicy, apply_param_dtype_map_patch + + apply_param_dtype_map_patch() + param_dtype_policy_cls = ParamDtypeMixedPrecisionPolicy + logger.info( + f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, " + f"param_dtype={param_dtype}, reduce_dtype={reduce_dtype}, " + f"param_dtype_overrides={param_dtype_maps.override_count} " + f"({param_dtype_maps.override_numel:,} parameters)" + ) + + fsdp_kwargs = { + "offload_policy": offload_policy, + "mesh": mesh, + } + + # input_dtype_policy owns boundary casts; autocast owns compute and keeps grad-ckpt recompute consistent. + def make_mp_policy(param_dtype_map): + if param_dtype_map: + assert param_dtype_policy_cls is not None + return param_dtype_policy_cls( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + cast_forward_inputs=False, + param_dtype_map=param_dtype_map, + ) + return MixedPrecisionPolicy( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + cast_forward_inputs=False, + ) + + for module, wrap_map in zip(modules, param_dtype_maps.wrap_maps, strict=True): + fully_shard( + module, + mp_policy=make_mp_policy(wrap_map), + **fsdp_kwargs, + ) + + fully_shard( + model, + mp_policy=make_mp_policy(param_dtype_maps.root_map), + **fsdp_kwargs, + ) + + return model diff --git a/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_integration_worker.py b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_integration_worker.py index 47a002003..1a4a52221 100644 --- a/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_integration_worker.py +++ b/tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_integration_worker.py @@ -37,7 +37,7 @@ import torch.distributed as dist from torch import nn -from miles.backends.fsdp_utils.actor import apply_fsdp2 +from miles.backends.fsdp_utils.model_loader import apply_fsdp2 from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan From 217276fc5df756972293282aa1d0b02474c5c31f Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 13:06:15 -0700 Subject: [PATCH 2/9] feat(fsdp): parameterize model loading for frozen roles --- miles/backends/fsdp_utils/actor.py | 10 ++- miles/backends/fsdp_utils/model_backend.py | 22 +++---- miles/backends/fsdp_utils/model_loader.py | 62 +++++++++---------- .../backends/fsdp_utils/models/ltx/loading.py | 4 +- miles/backends/fsdp_utils/models/package.py | 5 +- .../fsdp_utils/models/ltx/test_ltx.py | 32 +++++++--- .../backends/fsdp_utils/test_model_backend.py | 53 ++++++++++++++++ 7 files changed, 127 insertions(+), 61 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 1ec425904..1933d410d 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -93,7 +93,6 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty if self.args.start_rollout_id is None: self.args.start_rollout_id = 0 - self._master_dtype = parse_dtype_from_str(args.fsdp_master_dtype) self._forward_dtype = parse_dtype_from_str(args.diffusion_forward_dtype) from miles.utils.misc import load_function @@ -106,9 +105,14 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.model_backend.enable_deterministic_attention(args.fsdp_attention_backend) self.scheduler = self.model_backend.load_scheduler(args) self.models = load_fsdp_models( - args, self.model_backend, self.train_pipeline_config, self.parallel_state, master_dtype=self._master_dtype + args, + self.model_backend, + self.train_pipeline_config, + self.parallel_state, + checkpoint_path=args.hf_checkpoint, + trainable=True, + cpu_offload=args.fsdp_cpu_offload, ) - # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() clear_memory() diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index b1248a24d..b13277064 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -48,8 +48,8 @@ def enable_deterministic_attention(self, backend: str | None) -> None: def load_component( self, component: str, - args, *, + checkpoint_path: str, master_dtype: torch.dtype, materialize_weights: bool, ) -> torch.nn.Module: @@ -107,14 +107,14 @@ def enable_deterministic_attention(self, backend: str | None) -> None: def load_component( self, component: str, - args, *, + checkpoint_path: str, master_dtype: torch.dtype, materialize_weights: bool, ) -> torch.nn.Module: return self._pkg.loading.load_component( component, - args, + checkpoint_path=checkpoint_path, master_dtype=master_dtype, materialize_weights=materialize_weights, ) @@ -192,12 +192,12 @@ def _enable_deterministic_flash_attention(self) -> None: def load_component( self, component: str, - args, *, + checkpoint_path: str, master_dtype: torch.dtype, materialize_weights: bool, ) -> torch.nn.Module: - model_cls = self._resolve_component_class(args, component) + model_cls = self._resolve_component_class(checkpoint_path, component) kwargs = { "subfolder": component, "torch_dtype": master_dtype, @@ -212,17 +212,17 @@ def load_component( if not materialize_weights and keep_in_fp32 is not None: model_cls._keep_in_fp32_modules = None try: - return model_cls.from_pretrained(args.hf_checkpoint, **kwargs) + return model_cls.from_pretrained(checkpoint_path, **kwargs) finally: if not materialize_weights and keep_in_fp32 is not None: model_cls._keep_in_fp32_modules = keep_in_fp32 def load_scheduler(self, args) -> Any: - scheduler_cls = self._resolve_component_class(args, "scheduler") + scheduler_cls = self._resolve_component_class(args.hf_checkpoint, "scheduler") return scheduler_cls.from_pretrained(args.hf_checkpoint, subfolder="scheduler") @classmethod - def _resolve_component_class(cls, args, component: str): + def _resolve_component_class(cls, checkpoint_path: str, component: str): """Resolve ``component``'s class from ``model_index.json``. Components load individually via ``cls.from_pretrained(subfolder=...)`` rather @@ -233,13 +233,13 @@ def _resolve_component_class(cls, args, component: str): loaded from disk anyway — on every rank, and with ``low_cpu_mem_usage=False`` it also trips diffusers' ``_keep_in_fp32_modules`` guard. """ - config = DiffusionPipeline.load_config(args.hf_checkpoint) + config = DiffusionPipeline.load_config(checkpoint_path) if component not in config: - raise ValueError(f"pipeline {args.hf_checkpoint} has no component {component!r}") + raise ValueError(f"pipeline {checkpoint_path} has no component {component!r}") component_cls = cls._component_class(config[component]) if component_cls is None: raise ValueError( - f"cannot resolve the class for component {component!r} of {args.hf_checkpoint} " + f"cannot resolve the class for component {component!r} of {checkpoint_path} " f"from spec {config[component]!r}; remote-code components are not supported" ) return component_cls diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index 89cf4d019..0f07a42fa 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import logging import warnings from argparse import Namespace @@ -13,77 +15,71 @@ logger = logging.getLogger(__name__) -def load_fsdp_models(args: Namespace, model_backend, train_pipeline_config, parallel_state, *, master_dtype): - rank = dist.get_rank() - materialize_weights = rank == 0 - - models: dict[str, torch.nn.Module] = {} +def load_fsdp_models( + args: Namespace, + model_backend, + train_pipeline_config, + parallel_state, + *, + checkpoint_path: str, + trainable: bool = False, + cpu_offload: bool = False, +) -> dict[str, torch.nn.Module]: + materialize_weights = dist.get_rank() == 0 + master_dtype = parse_dtype_from_str(args.fsdp_master_dtype) + models = {} for component in args.update_weight_target_modules: - # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap with model_init_context(materialize_weights=materialize_weights): model = model_backend.load_component( component, - args, + checkpoint_path=checkpoint_path, master_dtype=master_dtype, materialize_weights=materialize_weights, ) if args.fsdp_attention_backend is not None: model_backend.set_attention_backend(model, args.fsdp_attention_backend) - - # Enable checkpointing on the raw model before PEFT wraps it. The flag - # is consumed when transformer blocks run, so LoRA layers inserted - # below remain inside the checkpointed block forward. - if args.gradient_checkpointing: + if trainable and args.gradient_checkpointing: model_backend.enable_gradient_checkpointing(model) - - if args.use_lora: + if trainable and args.use_lora: model = apply_lora(model, args, train_pipeline_config) + model.train(trainable) + if not trainable: + model.requires_grad_(False) - model.train() - - if rank != 0 and any(not parameter.is_meta for parameter in model.parameters()): + if not materialize_weights and any(not parameter.is_meta for parameter in model.parameters()): raise RuntimeError(f"{component} did not honor meta initialization") checkpoint.sync_model_dtypes(model) - full_state = model.state_dict() if rank == 0 else {} + full_state = model.state_dict() if materialize_weights else {} model = apply_fsdp2( model, model_backend.fsdp_parallel_plan(model), mesh=parallel_state.get_mesh("fsdp"), - cpu_offload=args.fsdp_cpu_offload, + cpu_offload=cpu_offload, args=args, ) - checkpoint.broadcast_full_state_to_fsdp( - model, - full_state, - cpu_offload=args.fsdp_cpu_offload, - ) + checkpoint.broadcast_full_state_to_fsdp(model, full_state, cpu_offload=cpu_offload) del full_state train_pipeline_config.postprocess_model_after_materialize(model) - models[component] = model - - if parallel_state.get_optional_mesh("sp") is not None: - for model in models.values(): - plan = model_backend.sequence_parallel_plan(model) + if parallel_state.get_optional_mesh("sp") is not None: apply_sequence_parallel( model, parallel_state, - plan, + model_backend.sequence_parallel_plan(model), model_backend.install_sequence_parallel_attention, ) + models[component] = model return models @contextmanager def model_init_context(*, materialize_weights: bool): - """Build real CPU weights on rank 0 and meta weights elsewhere.""" if materialize_weights: with torch.device("cpu"): yield return - from accelerate import init_empty_weights - # Some models compute buffer values during __init__, which cannot run on meta. + # Some constructors compute buffer values and cannot initialize them on meta. with init_empty_weights(include_buffers=False), warnings.catch_warnings(): warnings.filterwarnings( "ignore", diff --git a/miles/backends/fsdp_utils/models/ltx/loading.py b/miles/backends/fsdp_utils/models/ltx/loading.py index dbfefb34c..fc5c35a14 100644 --- a/miles/backends/fsdp_utils/models/ltx/loading.py +++ b/miles/backends/fsdp_utils/models/ltx/loading.py @@ -198,14 +198,14 @@ def resolve_transformer_checkpoint( def load_component( component: str, - args, *, + checkpoint_path: str, master_dtype: torch.dtype, materialize_weights: bool, ): if component != TRAIN_COMPONENT: raise ValueError(f"LTX trains the single DiT ({TRAIN_COMPONENT!r}); got {component!r}") - checkpoint = resolve_transformer_checkpoint(str(args.hf_checkpoint)) + checkpoint = resolve_transformer_checkpoint(checkpoint_path) return load_transformer_for_train( checkpoint, device="cpu", diff --git a/miles/backends/fsdp_utils/models/package.py b/miles/backends/fsdp_utils/models/package.py index fa123ecda..989ec2fe5 100644 --- a/miles/backends/fsdp_utils/models/package.py +++ b/miles/backends/fsdp_utils/models/package.py @@ -2,8 +2,9 @@ Each native model family is a Python package under ``models//`` with: - - ``loading`` — checkpoint resolution and ``load_component(..., - materialize_weights=...)``; distributed rank selection stays outside the package + - ``loading`` — checkpoint resolution and ``load_component(component, + checkpoint_path=..., master_dtype=..., materialize_weights=...)``; + distributed rank selection stays outside the package - ``modeling`` — ``load_scheduler``, ``enable_gradient_checkpointing``, optional ``flash_attention_entrypoints`` / ``required_flash_kernel_label`` - ``parallel_plan`` — ``FSDP_PARALLEL_PLAN``, ``sequence_parallel_plan`` diff --git a/tests/fast/backends/fsdp_utils/models/ltx/test_ltx.py b/tests/fast/backends/fsdp_utils/models/ltx/test_ltx.py index 469ad09fd..93de1f1c9 100644 --- a/tests/fast/backends/fsdp_utils/models/ltx/test_ltx.py +++ b/tests/fast/backends/fsdp_utils/models/ltx/test_ltx.py @@ -1,3 +1,10 @@ +"""Native LTX contracts: + +role checkpoint -> native resolver -> transformer loader +model family -> FSDP plan and attention backend +video geometry -> positions with matching latent token count +""" + from tests.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) @@ -25,29 +32,34 @@ def test_component_loading_delegates_to_loading_module(self, monkeypatch): sentinel = object() seen = {} - def fake_load_component(component, args, *, master_dtype, materialize_weights): + def resolve_checkpoint(checkpoint_path): + seen["checkpoint_path"] = checkpoint_path + return "resolved-transformer.safetensors" + + def load_transformer(checkpoint, *, device, dtype, materialize_weights): seen.update( - component=component, - args=args, - master_dtype=master_dtype, + checkpoint=checkpoint, + device=device, + dtype=dtype, materialize_weights=materialize_weights, ) return sentinel - monkeypatch.setattr(backend._pkg.loading, "load_component", fake_load_component) - args = object() + monkeypatch.setattr(backend._pkg.loading, "resolve_transformer_checkpoint", resolve_checkpoint) + monkeypatch.setattr(backend._pkg.loading, "load_transformer_for_train", load_transformer) result = backend.load_component( "transformer", - args, + checkpoint_path="teacher-checkpoint", master_dtype=torch.float32, materialize_weights=True, ) assert result is sentinel assert seen == { - "component": "transformer", - "args": args, - "master_dtype": torch.float32, + "checkpoint_path": "teacher-checkpoint", + "checkpoint": "resolved-transformer.safetensors", + "device": "cpu", + "dtype": torch.float32, "materialize_weights": True, } assert not hasattr(backend._pkg.modeling, "load_component") diff --git a/tests/fast/backends/fsdp_utils/test_model_backend.py b/tests/fast/backends/fsdp_utils/test_model_backend.py index 1c69122c8..c39ed2678 100644 --- a/tests/fast/backends/fsdp_utils/test_model_backend.py +++ b/tests/fast/backends/fsdp_utils/test_model_backend.py @@ -1,7 +1,16 @@ +"""Backend contracts: + +role checkpoint -> component loader; actor checkpoint -> scheduler +model family -> lifecycle hooks and FSDP parallel plan +""" + from tests.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) +import json +from argparse import Namespace + import torch from miles.backends.fsdp_utils.model_backend import BaseModelBackend, DiffusersModelBackend, MilesModelBackend @@ -22,6 +31,15 @@ def enable_gradient_checkpointing(self): self.gradient_checkpointing_enabled = True +class _CheckpointComponent(torch.nn.Module): + @classmethod + def from_pretrained(cls, checkpoint_path, **kwargs): + model = cls() + model.checkpoint_path = checkpoint_path + model.load_kwargs = kwargs + return model + + class TestBackendHierarchy: def test_concrete_backends_are_sibling_implementations(self): assert issubclass(MilesModelBackend, BaseModelBackend) @@ -45,3 +63,38 @@ def test_all_diffusers_model_plans_load(self): assert load_fsdp_parallel_plan("sd3").param_dtype_patterns == {} assert load_fsdp_parallel_plan("qwen_image").param_dtype_patterns == {} assert load_fsdp_parallel_plan("wan2_2").param_dtype_patterns + + def test_component_checkpoint_is_independent_of_actor_scheduler(self, tmp_path): + actor_checkpoint = tmp_path / "actor" + reference_checkpoint = tmp_path / "reference" + for checkpoint in (actor_checkpoint, reference_checkpoint): + checkpoint.mkdir() + (checkpoint / "model_index.json").write_text( + json.dumps( + { + "_class_name": "DiffusionPipeline", + "transformer": [__name__, "_CheckpointComponent"], + "scheduler": [__name__, "_CheckpointComponent"], + } + ) + ) + args = Namespace(hf_checkpoint=str(actor_checkpoint)) + backend = DiffusersModelBackend(None) + + model = backend.load_component( + "transformer", + checkpoint_path=str(reference_checkpoint), + master_dtype=torch.float32, + materialize_weights=True, + ) + scheduler = backend.load_scheduler(args) + + assert model.checkpoint_path == str(reference_checkpoint) + assert model.load_kwargs == { + "subfolder": "transformer", + "torch_dtype": torch.float32, + "low_cpu_mem_usage": True, + } + assert scheduler.checkpoint_path == str(actor_checkpoint) + assert scheduler.load_kwargs == {"subfolder": "scheduler"} + assert args == Namespace(hf_checkpoint=str(actor_checkpoint)) From b7875b2e181cda00c4484338426743b49ab1b43b Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 09:48:47 -0700 Subject: [PATCH 3/9] feat(fsdp): offload model and optimizer tensors to pinned CPU memory --- miles/backends/fsdp_utils/actor.py | 32 ++--- miles/backends/fsdp_utils/model_loader.py | 3 + miles/backends/fsdp_utils/offload.py | 71 ++++++++++ .../backends/fsdp_utils/_offload_worker.py | 124 ++++++++++++++++++ .../backends/fsdp_utils/test_offload.py | 41 ++++++ .../fast/backends/fsdp_utils/test_offload.py | 86 ++++++++++++ 6 files changed, 334 insertions(+), 23 deletions(-) create mode 100644 miles/backends/fsdp_utils/offload.py create mode 100644 tests/fast-gpu/backends/fsdp_utils/_offload_worker.py create mode 100644 tests/fast-gpu/backends/fsdp_utils/test_offload.py create mode 100644 tests/fast/backends/fsdp_utils/test_offload.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 1933d410d..f8b5d09eb 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -40,6 +40,7 @@ from .metrics import new_metric_buffer from .mixed_precision import parse_dtype_from_str from .model_loader import load_fsdp_models +from .offload import move_optimizer, offload_model, onload_model from .parallel import create_fsdp_parallel_state logger = logging.getLogger(__name__) @@ -84,9 +85,6 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty if self.args.debug_rollout_only: return 0 - if self.args.offload_train and self.args.fsdp_cpu_offload: - self.args.offload_train = False - if dist.get_rank() == 0: init_tracking(args, primary=False) @@ -195,9 +193,11 @@ def sleep(self) -> None: return print_memory("before offload DiT") + self.optimizer.zero_grad(set_to_none=True) - self.model.cpu() - move_torch_optimizer(self.optimizer, "cpu") + for model in self.models.values(): + offload_model(model) + move_optimizer(self.optimizer, "cpu") clear_memory() dist.barrier(group=get_gloo_group()) print_memory("after sleep DiT") @@ -207,8 +207,10 @@ def wake_up(self) -> None: if not self.args.offload_train: return - self.model.cuda() - move_torch_optimizer(self.optimizer, "cuda") + for model in self.models.values(): + onload_model(model, cpu_offload=self.args.fsdp_cpu_offload) + if not self.args.fsdp_cpu_offload: + move_optimizer(self.optimizer, "cuda") dist.barrier(group=get_gloo_group()) print_memory("after wake_up DiT") @@ -514,19 +516,3 @@ def _compute_noise_pred() -> torch.Tensor: write_old_log_prob=write_old_log_prob, old_log_prob_from_new=old_log_prob_from_new, ) - - -@torch.no_grad() -def move_torch_optimizer(optimizer, device): - """ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py""" - if not optimizer.state: - return - - for param_group in optimizer.param_groups: - for param in param_group["params"]: - state = optimizer.state[param] - for key, value in state.items(): - if isinstance(value, torch.Tensor): - state[key] = value.to(device, non_blocking=True) - - torch.cuda.synchronize() diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index 0f07a42fa..011a858e1 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -10,6 +10,7 @@ from . import checkpoint from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str +from .offload import offload_model from .sequence_parallel.plan import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -67,6 +68,8 @@ def load_fsdp_models( model_backend.sequence_parallel_plan(model), model_backend.install_sequence_parallel_attention, ) + if args.offload_train: + offload_model(model) models[component] = model return models diff --git a/miles/backends/fsdp_utils/offload.py b/miles/backends/fsdp_utils/offload.py new file mode 100644 index 000000000..45ed41ccf --- /dev/null +++ b/miles/backends/fsdp_utils/offload.py @@ -0,0 +1,71 @@ +import torch +from torch.distributed.fsdp import FSDPModule +from torch.utils._pytree import tree_map + + +@torch.no_grad() +def _move_model(model: torch.nn.Module, target_device: torch.device, *, keep_parameters_on_cpu: bool = False) -> None: + fsdp_parameters = [ + fsdp_parameter + for module in model.modules() + if isinstance(module, FSDPModule) + for fsdp_parameter_group in [module._get_fsdp_state()._fsdp_param_group] + if fsdp_parameter_group is not None + for fsdp_parameter in fsdp_parameter_group.fsdp_params + ] + # FSDP re-pads these uneven shards on the CPU as soon as _apply() hands them over. + repadded_parameter_ids = { + id(fsdp_parameter.sharded_param) + for fsdp_parameter in fsdp_parameters + if fsdp_parameter.sharded_param.to_local().size() != fsdp_parameter.padded_sharded_param_size + } + # _apply() visits a shared tensor once per registering module; reusing the first move keeps it shared. + original_and_moved_by_id = {} + + def move(tensor: torch.Tensor) -> torch.Tensor: + if id(tensor) not in original_and_moved_by_id: + if keep_parameters_on_cpu and isinstance(tensor, torch.nn.Parameter): + moved_tensor = tensor + else: + # A non_blocking device-to-host copy allocates its CPU output in pinned memory. + moved_tensor = tensor.to(target_device, non_blocking=True) + if id(tensor) in repadded_parameter_ids and tensor.is_cuda: + torch.cuda.current_stream().synchronize() + original_and_moved_by_id[id(tensor)] = (tensor, moved_tensor) + return original_and_moved_by_id[id(tensor)][1] + + original_pin_memory_by_fsdp_parameter = { + fsdp_parameter: fsdp_parameter.pin_memory for fsdp_parameter in fsdp_parameters + } + try: + if target_device.type == "cpu" and torch.cuda.is_available(): + # FSDP pins the storage it allocates for re-padding only when pin_memory is set. + for fsdp_parameter in fsdp_parameters: + fsdp_parameter.pin_memory = True + model._apply(move) + finally: + for fsdp_parameter, original_pin_memory in original_pin_memory_by_fsdp_parameter.items(): + fsdp_parameter.pin_memory = original_pin_memory + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def offload_model(model: torch.nn.Module) -> None: + """Move parameters and buffers to pinned CPU storage after gradients are cleared.""" + _move_model(model, torch.device("cpu")) + + +def onload_model(model: torch.nn.Module, *, cpu_offload: bool = False) -> None: + """Restore CUDA storage; native CPU offload keeps parameter shards on CPU.""" + _move_model(model, torch.device("cuda", torch.cuda.current_device()), keep_parameters_on_cpu=cpu_offload) + + +@torch.no_grad() +def move_optimizer(optimizer: torch.optim.Optimizer, device: str | torch.device) -> None: + # A non_blocking device-to-host copy allocates its CPU output in pinned memory. + for parameter, state in optimizer.state.items(): + optimizer.state[parameter] = tree_map( + lambda value: value.to(device, non_blocking=True) if isinstance(value, torch.Tensor) else value, state + ) + if torch.cuda.is_available(): + torch.cuda.synchronize() diff --git a/tests/fast-gpu/backends/fsdp_utils/_offload_worker.py b/tests/fast-gpu/backends/fsdp_utils/_offload_worker.py new file mode 100644 index 000000000..106df29fd --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/_offload_worker.py @@ -0,0 +1,124 @@ +"""Real CUDA counterpart to the CPU offload contract. + + same initial model --> control CUDA forward / AdamW + | + FSDP forward / AdamW --> pinned sleep --> wake --> equal next forward / AdamW + +The 2 x 2 x 2 matrix covers frozen references, cached gathered parameters, and +native CPUOffloadPolicy. Uneven local shards must remain pinned after FSDP padding; +optimizer Parameter bindings and tied/nonpersistent buffer aliases must survive. +""" + +import copy +import importlib.util +import itertools +import os +from pathlib import Path + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import CPUOffloadPolicy, OffloadPolicy, fully_shard +from torch.distributed.tensor import DTensor + +_MODULE_PATH = Path(__file__).resolve().parents[4] / "miles/backends/fsdp_utils/offload.py" +_SPEC = importlib.util.spec_from_file_location("fsdp_offload_under_test", _MODULE_PATH) +offload = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(offload) + + +class Tiny(torch.nn.Module): + def __init__(self): + super().__init__() + self.input = torch.nn.Linear(3, 5) + self.output = torch.nn.Linear(5, 3) + self.register_buffer("scale", torch.tensor([1.25])) + self.input.register_buffer("scale_alias", self.scale, persistent=False) + self.register_buffer("offset", torch.tensor([0.25]), persistent=False) + + def forward(self, inputs): + return self.output(self.input(inputs).tanh()) * self.scale + self.offset + + +def local_tensor(tensor): + return tensor.to_local() if isinstance(tensor, DTensor) else tensor + + +def check_model_offload_round_trip(mesh, frozen, reshard_after_forward, cpu_offload): + torch.manual_seed(19) + model = Tiny().cuda() + # Build aliases after cuda(), whose default _apply does not preserve buffer aliases. + model.input.scale_alias = model.scale + control = copy.deepcopy(model) + model.requires_grad_(not frozen) + control.requires_grad_(not frozen) + model.train(not frozen) + control.train(not frozen) + policy = CPUOffloadPolicy() if cpu_offload else OffloadPolicy() + fully_shard(model.input, mesh=mesh, reshard_after_forward=reshard_after_forward, offload_policy=policy) + fully_shard(model, mesh=mesh, reshard_after_forward=reshard_after_forward, offload_policy=policy) + parameters = dict(model.named_parameters()) + optimizer = None if frozen else torch.optim.AdamW(model.parameters(), lr=0.01) + control_optimizer = None if frozen else torch.optim.AdamW(control.parameters(), lr=0.01) + inputs = torch.arange(12, dtype=torch.float32, device="cuda").reshape(4, 3) / 12 + + for _ in range(2): + with torch.set_grad_enabled(not frozen): + actual = model(inputs) + expected = control(inputs) + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-6) + if not frozen: + actual.square().mean().backward() + expected.square().mean().backward() + optimizer.step() + control_optimizer.step() + optimizer.zero_grad(set_to_none=True) + control_optimizer.zero_grad(set_to_none=True) + offload.offload_model(model) + assert model.scale is model.input.scale_alias + assert "input.scale_alias" not in model.state_dict() + assert "offset" not in model.state_dict() + for name, parameter in model.named_parameters(): + assert parameter is parameters[name] + assert local_tensor(parameter).device.type == "cpu" and local_tensor(parameter).is_pinned(), name + assert all(buffer.device.type == "cpu" and buffer.is_pinned() for buffer in model.buffers()) + if optimizer is not None: + # Only state copied off the GPU lands in pinned memory; CPU-resident state (step, native offload) stays. + copied_from_cuda = { + (id(parameter), key): local_tensor(value).is_cuda + for parameter, state in optimizer.state.items() + for key, value in state.items() + } + offload.move_optimizer(optimizer, "cpu") + for parameter, state in optimizer.state.items(): + assert any(parameter is original for original in parameters.values()) + for key, value in state.items(): + assert local_tensor(value).device.type == "cpu" + assert local_tensor(value).is_pinned() or not copied_from_cuda[(id(parameter), key)], key + + offload.onload_model(model, cpu_offload=cpu_offload) + assert model.scale is model.input.scale_alias + expected_parameter_device = "cpu" if cpu_offload else "cuda" + assert all( + local_tensor(parameter).device.type == expected_parameter_device for parameter in model.parameters() + ) + assert all(buffer.device.type == "cuda" for buffer in model.buffers()) + if optimizer is not None: + offload.move_optimizer(optimizer, expected_parameter_device) + assert all(parameter is parameters[name] for name, parameter in model.named_parameters()) + + +def main(): + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl", device_id=torch.device("cuda", torch.cuda.current_device())) + mesh = init_device_mesh("cuda", (dist.get_world_size(),)) + for frozen, reshard_after_forward, cpu_offload in itertools.product((False, True), repeat=3): + check_model_offload_round_trip(mesh, frozen, reshard_after_forward, cpu_offload) + dist.barrier() + if dist.get_rank() == 0: + print("OK") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/fast-gpu/backends/fsdp_utils/test_offload.py b/tests/fast-gpu/backends/fsdp_utils/test_offload.py new file mode 100644 index 000000000..699aaaa3c --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/test_offload.py @@ -0,0 +1,41 @@ +"""Two-rank phase offload: FSDP shards + buffers + AdamW -> pinned CPU -> CUDA. + + trainable / frozen x reshard / no-reshard x native CPU offload / GPU + | + uneven shards + tied/nonpersistent buffers + | + offload --> pinned storage --> onload --> same outputs/updates +""" + +from tests.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, suite="stage-b-5-gpu-h200", labels=["fsdp"]) + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two CUDA devices") +def test_fsdp_phase_offload(): + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + "--nproc_per_node=2", + str(Path(__file__).with_name("_offload_worker.py")), + ], + env={**os.environ, "PYTHONUNBUFFERED": "1"}, + capture_output=True, + text=True, + timeout=300, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "OK" in result.stdout diff --git a/tests/fast/backends/fsdp_utils/test_offload.py b/tests/fast/backends/fsdp_utils/test_offload.py new file mode 100644 index 000000000..792b0c469 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_offload.py @@ -0,0 +1,86 @@ +"""Phase offload preserves training state and module tensor aliases. + + FSDP actor + AdamW --> clear gradients --> offload --> same Parameter objects + CPU state + | | + control step next actor step + +-------------------------- equal -------------------------+ + +CPU coverage checks FSDP bindings, cleared gradients, tied/nonpersistent buffers, and +optimizer state. The CUDA worker checks actual pinned allocations and transfers. +""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=["fsdp"]) + +import copy +import importlib.util +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.testing._internal.distributed.fake_pg # noqa: F401 +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import DTensor + +_MODULE_PATH = Path(__file__).resolve().parents[4] / "miles/backends/fsdp_utils/offload.py" +_SPEC = importlib.util.spec_from_file_location("fsdp_offload_under_test", _MODULE_PATH) +offload = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(offload) + + +@pytest.fixture +def cpu_mesh(): + dist.init_process_group("fake", store=dist.HashStore(), rank=0, world_size=1) + try: + yield init_device_mesh("cpu", (1,)) + finally: + dist.destroy_process_group() + + +def test_offload_preserves_training_bindings_and_buffer_aliases(cpu_mesh): + torch.manual_seed(9) + model = torch.nn.Sequential(torch.nn.Linear(3, 5), torch.nn.Tanh(), torch.nn.Linear(5, 3)).double() + model.register_buffer("marker", torch.tensor([7.0])) + model[0].register_buffer("marker_alias", model.marker, persistent=False) + control = copy.deepcopy(model) + fully_shard(model[0], mesh=cpu_mesh) + fully_shard(model, mesh=cpu_mesh) + parameters = dict(model.named_parameters()) + optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) + control_optimizer = torch.optim.AdamW(control.parameters(), lr=0.01) + inputs = torch.arange(12, dtype=torch.float64).reshape(4, 3) / 12 + + for _ in range(2): + for module, optim in ((model, optimizer), (control, control_optimizer)): + optim.zero_grad(set_to_none=True) + module(inputs).square().mean().backward() + optim.step() + optim.zero_grad(set_to_none=True) + offload.offload_model(model) + offload.move_optimizer(optimizer, "cpu") + assert model.marker is model[0].marker_alias + assert "0.marker_alias" not in model.state_dict() + for name, parameter in model.named_parameters(): + assert parameter is parameters[name] + assert parameter.grad is None + assert isinstance(parameter, DTensor) + torch.testing.assert_close(parameter.full_tensor(), control.get_parameter(name)) + assert all(parameter in optimizer.state for parameter in parameters.values()) + + +def test_optimizer_move_preserves_nested_state_and_metadata(): + parameter = torch.nn.Parameter(torch.tensor([1.0, 2.0])) + optimizer = torch.optim.SGD([parameter], lr=0.01) + optimizer.state[parameter] = { + "momentum_buffer": parameter.detach().clone(), + "schedule": {"step": 7, "history": [torch.tensor(3.0), "constant"]}, + } + offload.move_optimizer(optimizer, "cpu") + assert optimizer.param_groups[0]["params"][0] is parameter + assert optimizer.state[parameter]["schedule"]["step"] == 7 + assert optimizer.state[parameter]["schedule"]["history"][1] == "constant" + torch.testing.assert_close(optimizer.state[parameter]["momentum_buffer"], parameter) + torch.testing.assert_close(optimizer.state[parameter]["schedule"]["history"][0], torch.tensor(3.0)) From 6df88929c7212a4e9deed5b0465c775f18c147bd Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 09:48:52 -0700 Subject: [PATCH 4/9] feat(fsdp): load pretrained actor LoRA adapters and publish their metadata --- miles/backends/fsdp_utils/actor.py | 1 + .../diffusion_update_weight_utils.py | 6 +- miles/backends/fsdp_utils/model_loader.py | 18 +- miles/utils/arguments.py | 42 +++++ .../fsdp_utils/test_reference_models.py | 166 ++++++++++++++++++ tests/fast/utils/test_reference_arguments.py | 97 ++++++++++ 6 files changed, 327 insertions(+), 3 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/test_reference_models.py create mode 100644 tests/fast/utils/test_reference_arguments.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f8b5d09eb..cb31b8de0 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -108,6 +108,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.train_pipeline_config, self.parallel_state, checkpoint_path=args.hf_checkpoint, + lora_adapter_path=args.lora_adapter_path, trainable=True, cpu_offload=args.fsdp_cpu_offload, ) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index 0f68bf4d3..38e2670e0 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -295,9 +295,11 @@ def update_bucket_weights( "weight_version": str(weight_version), } if weight_update_mode is not None: + model = self.models[target_module] + adapter_config = model.peft_config[model.active_adapter] kwargs["weight_update_mode"] = weight_update_mode - kwargs["lora_alpha"] = self.args.lora_alpha - kwargs["lora_rank"] = self.args.lora_rank + kwargs["lora_alpha"] = adapter_config.lora_alpha + kwargs["lora_rank"] = adapter_config.r ref = self._ipc_engine.update_weights_from_tensor.remote(**kwargs) ray.get(ref) diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index 011a858e1..52c5a9775 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -23,6 +23,7 @@ def load_fsdp_models( parallel_state, *, checkpoint_path: str, + lora_adapter_path: str | None = None, trainable: bool = False, cpu_offload: bool = False, ) -> dict[str, torch.nn.Module]: @@ -41,7 +42,10 @@ def load_fsdp_models( model_backend.set_attention_backend(model, args.fsdp_attention_backend) if trainable and args.gradient_checkpointing: model_backend.enable_gradient_checkpointing(model) - if trainable and args.use_lora: + if lora_adapter_path is not None: + subfolder = component if len(args.update_weight_target_modules) > 1 else None + model = load_lora_adapter(model, lora_adapter_path, trainable=trainable, subfolder=subfolder) + elif trainable and args.use_lora: model = apply_lora(model, args, train_pipeline_config) model.train(trainable) if not trainable: @@ -91,6 +95,18 @@ def model_init_context(*, materialize_weights: bool): yield +def load_lora_adapter(model, adapter_path: str, *, trainable: bool, subfolder: str | None = None): + from peft import PeftConfig, PeftModel, get_peft_model + + if dist.get_rank() == 0: + return PeftModel.from_pretrained( + model, adapter_path, subfolder=subfolder, is_trainable=trainable, autocast_adapter_dtype=False + ) + config = PeftConfig.from_pretrained(adapter_path, subfolder=subfolder) + config.inference_mode = not trainable + return get_peft_model(model, config, low_cpu_mem_usage=True, autocast_adapter_dtype=False) + + def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -> torch.nn.Module: """Apply PEFT LoRA, leaving non-rank0 adapters uninitialized on meta.""" from peft import LoraConfig, get_peft_model diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index c8f7442a3..b2735441e 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1050,6 +1050,12 @@ def add_lora_arguments(parser): parser.add_argument( "--use-lora", action="store_true", default=False, help="Use LoRA adapters instead of full finetune." ) + parser.add_argument( + "--lora-adapter-path", + type=str, + default=None, + help="Pretrained actor LoRA adapter to continue training.", + ) parser.add_argument("--lora-rank", type=int, default=64) parser.add_argument("--lora-alpha", type=int, default=64) parser.add_argument( @@ -1571,6 +1577,41 @@ def set_default_diffusion_args(args) -> None: args.ref_mode = "none" +def validate_actor_lora_adapter(args) -> None: + if args.lora_adapter_path is None: + return + if not args.use_lora: + raise ValueError("--lora-adapter-path requires --use-lora") + from peft import PeftConfig + + for component in args.update_weight_target_modules: + subfolder = component if len(args.update_weight_target_modules) > 1 else None + adapter_config = PeftConfig.from_pretrained(args.lora_adapter_path, subfolder=subfolder) + if ( + adapter_config.use_dora + or adapter_config.bias != "none" + or adapter_config.lora_bias + or adapter_config.modules_to_save + ): + raise ValueError( + f"Actor LoRA adapter for {component} must contain only LoRA A/B matrices; " + "DoRA, bias training, and modules_to_save are not supported by checkpoint and weight sync" + ) + if args.lora_ipc_weight_sync and ( + adapter_config.use_rslora or adapter_config.rank_pattern or adapter_config.alpha_pattern + ): + raise ValueError( + f"--lora-ipc-weight-sync requires uniform standard LoRA for {component}; " + "disable it for rsLoRA or per-layer rank/alpha" + ) + if (adapter_config.r, adapter_config.lora_alpha) != (args.lora_rank, args.lora_alpha): + raise ValueError( + f"Actor LoRA adapter for {component} has r={adapter_config.r}, lora_alpha={adapter_config.lora_alpha}; " + f"set --lora-rank {adapter_config.r} --lora-alpha {adapter_config.lora_alpha} " + f"instead of {args.lora_rank}/{args.lora_alpha}" + ) + + def miles_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args) @@ -1608,6 +1649,7 @@ def miles_validate_args(args): ) if len(set(args.update_weight_target_modules)) != len(args.update_weight_target_modules): raise ValueError(f"--update-weight-target-module has duplicates: {args.update_weight_target_module!r}") + validate_actor_lora_adapter(args) if args.wandb_log_image_interval < 1: raise ValueError(f"wandb_log_image_interval must be >= 1, got {args.wandb_log_image_interval}") diff --git a/tests/fast/backends/fsdp_utils/test_reference_models.py b/tests/fast/backends/fsdp_utils/test_reference_models.py new file mode 100644 index 000000000..efa7cb8b2 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_reference_models.py @@ -0,0 +1,166 @@ +"""Loaded actor adapters drive IPC metadata instead of CLI initialization defaults. + + component checkpoints --> actual PEFT loader --> loaded r/alpha --> captured IPC payload + rank 0 / unsharded CPU fixture --> no Ray or CUDA transport required + checkpoint_path --> backend loads base; loader applies LoRA +""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=20, suite="stage-a-cpu", labels=["fsdp"]) + +import ast +import logging +from argparse import Namespace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from peft import LoraConfig, get_peft_model + +from miles.backends.fsdp_utils import model_loader +from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan + + +class TinyBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(3, 4, bias=False) + + def forward(self, inputs): + return self.proj(inputs).tanh() + + +class TinyComponent(torch.nn.Module): + def __init__(self): + super().__init__() + self.block = TinyBlock() + self.register_buffer("forward_count", torch.zeros((), dtype=torch.int64), persistent=False) + + def forward(self, inputs): + self.forward_count.add_(1) + return self.block(inputs) + + +class TinyBackend: + def __init__(self): + self.component_load_requests = [] + + def load_component(self, component, *, checkpoint_path, master_dtype, materialize_weights): + self.component_load_requests.append((component, checkpoint_path)) + model = TinyComponent().to(dtype=master_dtype) + if materialize_weights: + value = {"actor": 0.1, "reference": 0.2, "teacher": 0.3}[checkpoint_path] + value += 0.05 if component == "transformer_2" else 0.0 + with torch.no_grad(): + model.block.proj.weight.fill_(value) + return model + + def fsdp_parallel_plan(self, model): + return FSDPParallelPlan(no_split_modules=("TinyBlock",)) + + +class TinyPipelineConfig: + input_dtype_policy = {} + + def postprocess_model_after_materialize(self, model): + pass + + def compute_noise_pred(self, *, model, latents_input, **kwargs): + return model(latents_input) + + +def make_model_loader_args(): + return Namespace( + hf_checkpoint="actor", + use_lora=True, + update_weight_target_modules=["transformer", "transformer_2"], + fsdp_master_dtype="fp32", + fsdp_reduce_dtype="fp32", + diffusion_forward_dtype="fp32", + fsdp_attention_backend=None, + gradient_checkpointing=False, + offload_train=False, + ) + + +def save_component_lora_adapters(path): + for index, component in enumerate(make_model_loader_args().update_weight_target_modules): + model = get_peft_model(TinyComponent(), LoraConfig(r=2, lora_alpha=4, target_modules=["proj"])) + with torch.no_grad(): + model.base_model.model.block.proj.lora_A.default.weight.fill_(0.15 + 0.1 * index) + model.base_model.model.block.proj.lora_B.default.weight.fill_(0.25 + 0.1 * index) + model.save_pretrained(path / component) + + +@pytest.fixture +def unsharded_model_loader(monkeypatch): + calls = [] + monkeypatch.setattr(model_loader.dist, "get_rank", lambda: 0) + monkeypatch.setattr(model_loader.checkpoint, "sync_model_dtypes", lambda model: None) + monkeypatch.setattr(model_loader.checkpoint, "broadcast_full_state_to_fsdp", lambda *args, **kwargs: None) + + def record_fsdp(model, plan, *, mesh, cpu_offload, args): + calls.append((mesh, cpu_offload)) + return model + + monkeypatch.setattr(model_loader, "apply_fsdp2", record_fsdp) + parallel = SimpleNamespace(get_mesh=lambda name: "shared-fsdp-mesh", get_optional_mesh=lambda name: None) + return parallel, calls + + +def test_ipc_publication_uses_loaded_adapter_config(tmp_path, unsharded_model_loader): + parallel, _ = unsharded_model_loader + save_component_lora_adapters(tmp_path) + models = model_loader.load_fsdp_models( + make_model_loader_args(), + TinyBackend(), + TinyPipelineConfig(), + parallel, + checkpoint_path="actor", + lora_adapter_path=str(tmp_path), + trainable=True, + ) + path = Path(__file__).resolve().parents[4] / "miles/backends/fsdp_utils/diffusion_update_weight_utils.py" + source = ast.parse(path.read_text()) + updater = next( + node + for node in source.body + if isinstance(node, ast.ClassDef) and node.name == "DiffusionUpdateWeightFromTensor" + ) + method = next( + node for node in updater.body if isinstance(node, ast.FunctionDef) and node.name == "update_bucket_weights" + ) + sent = [] + + def gather_object(*, obj, object_gather_list, **kwargs): + object_gather_list[0] = obj + + namespace = dict( + logger=logging.getLogger(__name__), + monkey_patch_torch_reductions=lambda: None, + FlattenedTensorBucket=lambda named_tensors: SimpleNamespace( + get_metadata=lambda: [], + get_flattened_tensor=lambda: named_tensors[0][1], + ), + MultiprocessingSerializer=SimpleNamespace(serialize=lambda *args, **kwargs: "payload"), + dist=SimpleNamespace(get_rank=lambda: 0, get_world_size=lambda group: 1, gather_object=gather_object), + get_physical_gpu_id=lambda: "gpu-0", + ray=SimpleNamespace(get=lambda result: result), + ) + exec(compile(ast.Module(body=[method], type_ignores=[]), str(path), "exec"), namespace) + harness = SimpleNamespace( + args=Namespace(lora_rank=64, lora_alpha=64), + models=models, + _ipc_gather_src=0, + _ipc_gather_group=None, + _ipc_engine=SimpleNamespace( + update_weights_from_tensor=SimpleNamespace(remote=lambda **kwargs: sent.append(kwargs)) + ), + ) + namespace[method.name]( + harness, [("block.proj.lora_A.weight", torch.ones(2, 3))], "transformer", weight_update_mode="lora_merge" + ) + assert len(sent) == 1 + assert sent[0]["lora_rank"] == 2 and sent[0]["lora_alpha"] == 4 diff --git a/tests/fast/utils/test_reference_arguments.py b/tests/fast/utils/test_reference_arguments.py new file mode 100644 index 000000000..1bb52a3a0 --- /dev/null +++ b/tests/fast/utils/test_reference_arguments.py @@ -0,0 +1,97 @@ +"""Pretrained actor LoRA must match checkpoint and weight-publication support. + + --use-lora --> saved adapter config --> A/B-only checkpoint + uniform IPC scaling + CLI r/alpha match + unsupported config -----------------> reject before model construction + +Validation requires the CLI rank/alpha to match the loaded adapter and leaves args unchanged. +""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=["fsdp"]) + +from argparse import Namespace + +import pytest +from peft import LoraConfig + +from miles.utils.arguments import validate_actor_lora_adapter + + +def _args(**overrides): + values = dict( + ref_mode="none", + ref_load=None, + teacher_load=None, + lora_adapter_path=None, + lora_ipc_weight_sync=False, + lora_rank=2, + lora_alpha=4, + update_weight_target_modules=["transformer", "transformer_2"], + ref_lora_adapter_path=None, + teacher_lora_adapter_path=None, + ref_cpu_offload=False, + teacher_cpu_offload=False, + use_lora=False, + use_ema=False, + sglang_tp_size=1, + sglang_sp_degree=None, + sglang_enable_cfg_parallel=False, + loss_type="policy_loss", + diffusion_kl_beta=0.0, + ) + values.update(overrides) + return Namespace(**values) + + +def test_loaded_adapter_matching_actor_cli_passes_without_changing_args(tmp_path): + # adapter r=2 alpha=4 == --lora-rank 2 --lora-alpha 4 -> accepted, args untouched + args = _args(use_lora=True, lora_adapter_path=str(tmp_path), lora_ipc_weight_sync=True) + for component in args.update_weight_target_modules: + LoraConfig(r=2, lora_alpha=4, target_modules=["proj"]).save_pretrained(tmp_path / component) + original = vars(args).copy() + validate_actor_lora_adapter(args) + assert vars(args) == original + + +@pytest.mark.parametrize("cli", [{"lora_rank": 64}, {"lora_alpha": 64}]) +def test_loaded_adapter_rejects_mismatched_actor_cli(tmp_path, cli): + # adapter r=2 alpha=4 vs CLI 64 -> rejected instead of silently ignoring the CLI values + LoraConfig(r=2, lora_alpha=4, target_modules=["proj"]).save_pretrained(tmp_path) + args = _args(use_lora=True, lora_adapter_path=str(tmp_path), update_weight_target_modules=["transformer"], **cli) + with pytest.raises(ValueError, match="set --lora-rank 2 --lora-alpha 4"): + validate_actor_lora_adapter(args) + + +def test_actor_adapter_requires_lora_training_before_loading_config(): + args = _args(lora_adapter_path="unread-adapter-checkpoint") + with pytest.raises(ValueError, match="--lora-adapter-path requires --use-lora"): + validate_actor_lora_adapter(args) + + +@pytest.mark.parametrize( + ("config_options", "ipc", "message"), + [ + ({"use_dora": True}, False, "must contain only LoRA A/B matrices"), + ({"use_rslora": True}, True, "requires uniform standard LoRA"), + ], +) +def test_loaded_adapter_rejects_unsupported_export_state(tmp_path, config_options, ipc, message): + LoraConfig(r=2, lora_alpha=4, target_modules=["proj"], **config_options).save_pretrained(tmp_path) + args = _args( + use_lora=True, + lora_adapter_path=str(tmp_path), + lora_ipc_weight_sync=ipc, + update_weight_target_modules=["transformer"], + ) + with pytest.raises(ValueError, match=message): + validate_actor_lora_adapter(args) + + +def test_layer_specific_scaling_is_allowed_without_ipc(tmp_path): + LoraConfig( + r=2, lora_alpha=4, use_rslora=True, rank_pattern={"proj": 4}, alpha_pattern={"proj": 8} + ).save_pretrained(tmp_path) + validate_actor_lora_adapter( + _args(use_lora=True, lora_adapter_path=str(tmp_path), update_weight_target_modules=["transformer"]) + ) From 4f09c063ce42ded75ece1728da6798365a5480c7 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 09:48:57 -0700 Subject: [PATCH 5/9] fix(fsdp): isolate base-reference inference from actor training state --- miles/backends/fsdp_utils/actor.py | 33 ++-- miles/backends/fsdp_utils/ema.py | 8 + miles/backends/fsdp_utils/model_loader.py | 10 ++ .../fsdp_utils/test_reference_models.py | 149 +++++++++++++++++- 4 files changed, 179 insertions(+), 21 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index cb31b8de0..b739535be 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -33,7 +33,7 @@ DiffusionUpdateWeightFromTensorLoRA, DiffusionUpdateWeightFromTensorLoRAIPC, ) -from .ema import EmaShadow +from .ema import EmaShadow, reshard_model from .input_dtype_policy import apply_input_dtype_policy from .loss_hub import DiffusionLossContext, flow_grpo_loss_formula, prepare_flow_grpo_batch from .lr_scheduler import get_lr_scheduler @@ -295,12 +295,6 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: num_pairs = len(train_pairs) - ref_mode = self.args.ref_mode - if ref_mode == "lora_base" and not all(hasattr(m, "disable_adapter") for m in self.models.values()): - raise RuntimeError( - "--ref-mode lora_base requires PEFT models exposing disable_adapter() after FSDP wrapping." - ) - # ------------- Rollout Scheduler Metadata ------------- scheduler_timesteps, scheduler_sigmas = scheduler_meta_from_rollout( rollout_data, @@ -469,10 +463,10 @@ def _forward_train_pair_batch( default_dtype=forward_dtype, ) - def _compute_noise_pred() -> torch.Tensor: + def _compute_noise_pred(model) -> torch.Tensor: with torch.autocast("cuda", dtype=forward_dtype, enabled=forward_dtype != torch.float32): return train_pipeline_config.compute_noise_pred( - model=prepared.model, + model=model, latents_input=latents_in, timesteps_input=timesteps_in, pos_cond=pos_cond_in, @@ -484,17 +478,20 @@ def _compute_noise_pred() -> torch.Tensor: true_cfg_scale=prepared.true_cfg_scale, ) - new_pred = _compute_noise_pred() + new_pred = _compute_noise_pred(prepared.model) ref_pred = None - ref_mode = self.args.ref_mode - if ref_mode != "none": - if ref_mode == "ema": - ref_ctx = self.ema_shadow.swap_in() - else: - ref_ctx = prepared.model.disable_adapter() - with torch.no_grad(), ref_ctx: - ref_pred = _compute_noise_pred().detach() + if self.args.ref_mode == "ema": + # Drop gathered copies around the swap so neither the EMA nor the actor forward sees stale weights. + reshard_model(prepared.model) + with torch.no_grad(), self.ema_shadow.swap_in(): + ref_pred = _compute_noise_pred(prepared.model).detach() + reshard_model(prepared.model) + elif self.args.ref_mode == "lora_base": + with torch.no_grad(), prepared.model.disable_adapter(): + ref_pred = _compute_noise_pred(prepared.model).detach() + # PEFT re-enables adapter gradients on exit; reshard first so that lands on the shards, not gathered copies. + reshard_model(prepared.model) if self.custom_loss_formula_func is not None: return self.custom_loss_formula_func( diff --git a/miles/backends/fsdp_utils/ema.py b/miles/backends/fsdp_utils/ema.py index 9120ef3b8..aef9d0711 100644 --- a/miles/backends/fsdp_utils/ema.py +++ b/miles/backends/fsdp_utils/ema.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from torch.distributed.fsdp import FSDPModule from torch.distributed.tensor import DTensor @@ -14,6 +15,13 @@ def _local(t: torch.Tensor) -> torch.Tensor: return t.to_local() if isinstance(t, DTensor) else t +def reshard_model(model: torch.nn.Module) -> None: + """Drop FSDP's gathered full parameters so the registered parameters are the local shards again.""" + for module in model.modules(): + if isinstance(module, FSDPModule): + module.reshard() + + class EmaShadow: """EMA shadow of trainable parameters.""" diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index 52c5a9775..3d8e1d581 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -72,12 +72,22 @@ def load_fsdp_models( model_backend.sequence_parallel_plan(model), model_backend.install_sequence_parallel_attention, ) + finish_fsdp_lazy_init(model) if args.offload_train: offload_model(model) models[component] = model return models +def finish_fsdp_lazy_init(model: torch.nn.Module) -> None: + """Run FSDP's one-time lazy init now instead of at the first forward. + + It fixes each group's gradient dtype from the parameters that require grad at that moment; + a first forward with PEFT adapters disabled would see nothing trainable and leave it unset. + """ + model._get_fsdp_state()._lazy_init() + + @contextmanager def model_init_context(*, materialize_weights: bool): if materialize_weights: diff --git a/tests/fast/backends/fsdp_utils/test_reference_models.py b/tests/fast/backends/fsdp_utils/test_reference_models.py index efa7cb8b2..8299a2620 100644 --- a/tests/fast/backends/fsdp_utils/test_reference_models.py +++ b/tests/fast/backends/fsdp_utils/test_reference_models.py @@ -1,7 +1,9 @@ -"""Loaded actor adapters drive IPC metadata instead of CLI initialization defaults. +"""Loaded actor adapters preserve base-reference training and publication metadata. - component checkpoints --> actual PEFT loader --> loaded r/alpha --> captured IPC payload - rank 0 / unsharded CPU fixture --> no Ray or CUDA transport required + checkpoint --> component LoRA weights --> base reference --> actor gradients + loaded r/alpha -----------------------> IPC publication (ignores CLI defaults) + +Real PEFT loading and actor forward are exercised with FSDP collectives replaced. checkpoint_path --> backend loads base; loader applies LoRA """ @@ -10,6 +12,7 @@ register_cpu_ci(est_time=20, suite="stage-a-cpu", labels=["fsdp"]) import ast +import copy import logging from argparse import Namespace from pathlib import Path @@ -17,9 +20,15 @@ import pytest import torch +import torch.distributed as dist +import torch.testing._internal.distributed.fake_pg # noqa: F401 from peft import LoraConfig, get_peft_model +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard from miles.backends.fsdp_utils import model_loader +from miles.backends.fsdp_utils.ema import reshard_model +from miles.backends.fsdp_utils.input_dtype_policy import apply_input_dtype_policy from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan @@ -94,6 +103,57 @@ def save_component_lora_adapters(path): model.save_pretrained(path / component) +def load_actor_forward_method(): + path = Path(__file__).resolve().parents[4] / "miles/backends/fsdp_utils/actor.py" + source = ast.parse(path.read_text()) + actor = next(node for node in source.body if isinstance(node, ast.ClassDef) and node.name == "FSDPTrainRayActor") + method = next( + node for node in actor.body if isinstance(node, ast.FunctionDef) and node.name == "_forward_train_pair_batch" + ) + namespace = dict( + torch=torch, + reshard_model=reshard_model, + apply_input_dtype_policy=apply_input_dtype_policy, + DiffusionLossContext=object, + MetricBuffer=object, + ) + exec(compile(ast.Module(body=[method], type_ignores=[]), str(path), "exec"), namespace) + return namespace[method.name] + + +def make_actor_forward_harness(models, reference_models, component, inputs, ref_mode): + observed = {} + prepared = SimpleNamespace( + model=models[component], + component_name=component, + latents=inputs, + timesteps_for_model=torch.ones(inputs.shape[0], device=inputs.device), + pos_cond=None, + neg_cond=None, + joint_cond=None, + use_cfg=False, + cfg_batching=False, + guidance_scale=1.0, + true_cfg_scale=None, + ) + + def loss_formula(ctx, batch, prepared, *, new_pred, ref_pred, **kwargs): + observed.update(new_pred=new_pred, ref_pred=ref_pred) + return (new_pred - ref_pred).square().mean() + + harness = SimpleNamespace( + args=Namespace(ref_mode=ref_mode), + models=models, + model=torch.nn.ModuleDict(models), + reference_models=reference_models, + train_pipeline_config=TinyPipelineConfig(), + _forward_dtype=torch.float32, + custom_prepare_train_batch_func=lambda *args, **kwargs: prepared, + custom_loss_formula_func=loss_formula, + ) + return harness, observed + + @pytest.fixture def unsharded_model_loader(monkeypatch): calls = [] @@ -106,10 +166,93 @@ def record_fsdp(model, plan, *, mesh, cpu_offload, args): return model monkeypatch.setattr(model_loader, "apply_fsdp2", record_fsdp) + monkeypatch.setattr(model_loader, "finish_fsdp_lazy_init", lambda model: None) parallel = SimpleNamespace(get_mesh=lambda name: "shared-fsdp-mesh", get_optional_mesh=lambda name: None) return parallel, calls +def test_lora_base_forward_restores_adapters_and_actor_gradients(tmp_path, unsharded_model_loader): + parallel, _ = unsharded_model_loader + save_component_lora_adapters(tmp_path) + actor = model_loader.load_fsdp_models( + make_model_loader_args(), + TinyBackend(), + TinyPipelineConfig(), + parallel, + checkpoint_path="actor", + lora_adapter_path=str(tmp_path), + trainable=True, + ) + component = "transformer_2" + model, inputs = actor[component], torch.ones(2, 3) + oracle = copy.deepcopy(model) + base = TinyComponent() + base.block.proj.weight.data.copy_(model.base_model.model.block.proj.base_layer.weight) + with torch.no_grad(): + target = base(inputs) + harness, observed = make_actor_forward_harness(actor, {}, component, inputs, "lora_base") + load_actor_forward_method()(harness, None, [], metrics=None).backward() + oracle_output = oracle(inputs) + (oracle_output - target).square().mean().backward() + torch.testing.assert_close(observed["ref_pred"], target) + torch.testing.assert_close(observed["new_pred"], oracle_output) + assert not torch.equal(target, oracle_output) + assert model.training and not model.base_model.model.block.proj.disable_adapters + assert model.base_model.model.forward_count.item() == 2 # reference forward + actor forward + assert actor["transformer"].base_model.model.forward_count.item() == 0 + for actual, expected in zip(model.parameters(), oracle.parameters(), strict=True): + if actual.requires_grad: + torch.testing.assert_close(actual.grad, expected.grad) + + +def test_first_forward_with_adapters_disabled_keeps_master_dtype_gradients(): + # FSDP2 fixes each group's gradient dtype at its lazy init from the parameters that require grad: + # load --finish_fsdp_lazy_init (adapters trainable)--> lora_base forward FIRST --> actor backward + # --> fp32 adapter gradients on the fp32 shards, whichever forward runs first + dist.init_process_group("fake", store=dist.HashStore(), rank=0, world_size=1) + try: + mesh = init_device_mesh("cpu", (1,)) + model = get_peft_model(TinyComponent(), LoraConfig(r=2, lora_alpha=4, target_modules=["proj"])) + policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) + fully_shard(model.base_model.model.block, mesh=mesh, mp_policy=policy) + fully_shard(model, mesh=mesh, mp_policy=policy) + model_loader.finish_fsdp_lazy_init(model) + inputs = torch.ones(2, 3) + with torch.no_grad(), model.disable_adapter(): + model(inputs) + reshard_model(model) + model(inputs).square().mean().backward() + for name, parameter in model.named_parameters(): + if "lora_" in name: + assert parameter.grad is not None and parameter.grad.dtype == torch.float32, name + finally: + dist.destroy_process_group() + + +def test_lora_base_reference_keeps_adapter_gradients_on_fsdp_shards(): + # FSDP2 with reshard_after_forward=False keeps the gathered parameters registered after forward: + # lora_base forward --PEFT re-enables adapter gradients on exit--> must land on the shards + # --> every following training step still produces adapter gradients + dist.init_process_group("fake", store=dist.HashStore(), rank=0, world_size=1) + try: + mesh = init_device_mesh("cpu", (1,)) + model = get_peft_model(TinyComponent(), LoraConfig(r=2, lora_alpha=4, target_modules=["proj"])) + fully_shard(model.base_model.model.block, mesh=mesh, reshard_after_forward=False) + fully_shard(model, mesh=mesh, reshard_after_forward=False) + adapter_parameters = {name: parameter for name, parameter in model.named_parameters() if "lora_" in name} + forward = load_actor_forward_method() + for step in range(2): + harness, _ = make_actor_forward_harness( + {"transformer": model}, {}, "transformer", torch.ones(2, 3), "lora_base" + ) + forward(harness, None, [], metrics=None).backward() + for name, parameter in adapter_parameters.items(): + assert parameter.requires_grad and parameter.grad is not None, (step, name) + model.zero_grad(set_to_none=True) + finally: + dist.destroy_process_group() + + def test_ipc_publication_uses_loaded_adapter_config(tmp_path, unsharded_model_loader): parallel, _ = unsharded_model_loader save_component_lora_adapters(tmp_path) From df445088790375d444d78b138de24da280a8fb6c Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 09:49:02 -0700 Subject: [PATCH 6/9] feat(fsdp): load independent frozen reference and teacher models --- miles/backends/fsdp_utils/actor.py | 44 ++++- .../backends/fsdp_utils/loss_hub/flow_grpo.py | 2 +- miles/backends/fsdp_utils/loss_hub/types.py | 2 + miles/utils/arguments.py | 59 +++++-- .../fsdp_utils/_reference_models_worker.py | 151 ++++++++++++++++++ .../fsdp_utils/test_reference_models.py | 43 +++++ .../fsdp_utils/test_reference_models.py | 112 ++++++++++++- tests/fast/utils/test_reference_arguments.py | 57 ++++++- 8 files changed, 442 insertions(+), 28 deletions(-) create mode 100644 tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py create mode 100644 tests/fast-gpu/backends/fsdp_utils/test_reference_models.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index b739535be..192233239 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -112,6 +112,29 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty trainable=True, cpu_offload=args.fsdp_cpu_offload, ) + self.reference_models = {} + if args.ref_load is not None: + self.reference_models = load_fsdp_models( + args, + self.model_backend, + self.train_pipeline_config, + self.parallel_state, + checkpoint_path=args.ref_load, + lora_adapter_path=args.ref_lora_adapter_path, + cpu_offload=args.ref_cpu_offload, + ) + self.teacher_models = {} + if args.teacher_load is not None: + self.teacher_models = load_fsdp_models( + args, + self.model_backend, + self.train_pipeline_config, + self.parallel_state, + checkpoint_path=args.teacher_load, + lora_adapter_path=args.teacher_lora_adapter_path, + cpu_offload=args.teacher_cpu_offload, + ) + # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() clear_memory() @@ -196,8 +219,9 @@ def sleep(self) -> None: print_memory("before offload DiT") self.optimizer.zero_grad(set_to_none=True) - for model in self.models.values(): - offload_model(model) + for models in (self.models, self.reference_models, self.teacher_models): + for model in models.values(): + offload_model(model) move_optimizer(self.optimizer, "cpu") clear_memory() dist.barrier(group=get_gloo_group()) @@ -208,8 +232,13 @@ def wake_up(self) -> None: if not self.args.offload_train: return - for model in self.models.values(): - onload_model(model, cpu_offload=self.args.fsdp_cpu_offload) + for models, cpu_offload in ( + (self.models, self.args.fsdp_cpu_offload), + (self.reference_models, self.args.ref_cpu_offload), + (self.teacher_models, self.args.teacher_cpu_offload), + ): + for model in models.values(): + onload_model(model, cpu_offload=cpu_offload) if not self.args.fsdp_cpu_offload: move_optimizer(self.optimizer, "cuda") dist.barrier(group=get_gloo_group()) @@ -332,6 +361,8 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: loss_ctx = DiffusionLossContext( models=self.models, + reference_models=self.reference_models, + teacher_models=self.teacher_models, train_pipeline_config=self.train_pipeline_config, sde_backend=self.sde_backend, scheduler=self.scheduler, @@ -481,7 +512,10 @@ def _compute_noise_pred(model) -> torch.Tensor: new_pred = _compute_noise_pred(prepared.model) ref_pred = None - if self.args.ref_mode == "ema": + if self.args.ref_mode == "ref": + with torch.no_grad(): + ref_pred = _compute_noise_pred(self.reference_models[prepared.component_name]).detach() + elif self.args.ref_mode == "ema": # Drop gathered copies around the swap so neither the EMA nor the actor forward sees stale weights. reshard_model(prepared.model) with torch.no_grad(), self.ema_shadow.swap_in(): diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py index ee49bbd1b..5893adcd9 100644 --- a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py +++ b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py @@ -149,7 +149,7 @@ def flow_grpo_loss_formula( kl_sum = loss_sum.new_zeros(()) if kl_beta > 0: if ref_pred is None: - raise ValueError("Flow-GRPO KL requires a reference DiT forward (actor ref_mode=lora_base)") + raise ValueError("Flow-GRPO KL requires a reference DiT forward; set --ref-mode lora_base, ema, or ref") _, _, prev_sample_mean_ref, _ = ctx.sde_backend.sde_step_logprob( ref_pred.float(), prepared.timesteps, diff --git a/miles/backends/fsdp_utils/loss_hub/types.py b/miles/backends/fsdp_utils/loss_hub/types.py index 74d252b44..d1e3d831b 100644 --- a/miles/backends/fsdp_utils/loss_hub/types.py +++ b/miles/backends/fsdp_utils/loss_hub/types.py @@ -24,6 +24,8 @@ class DiffusionLossContext: rollout_id: int = 0 microbatch_id: int = 0 dp_rank: int = 0 + reference_models: dict[str, torch.nn.Module] = field(default_factory=dict) + teacher_models: dict[str, torch.nn.Module] = field(default_factory=dict) @dataclass diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index b2735441e..7ce2b0ca9 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -962,14 +962,39 @@ def add_algo_arguments(parser): parser.add_argument( "--ref-mode", type=str, - choices=["none", "lora_base", "ema"], + choices=["none", "lora_base", "ema", "ref"], default=None, help=( - "Which reference weights to use for the no-grad DiT forward. " + "Which reference model to use for the no-grad DiT forward. " "Auto: lora_base when --diffusion-kl-beta > 0 and ema for --loss-type nft. " - "Explicit values skip auto inference." + "Use ref with --ref-load for an independent frozen model." ), ) + parser.add_argument( + "--ref-load", type=str, default=None, help="HF pipeline checkpoint for the frozen reference model." + ) + parser.add_argument( + "--teacher-load", type=str, default=None, help="HF pipeline checkpoint for the frozen teacher model." + ) + parser.add_argument( + "--ref-lora-adapter-path", type=str, default=None, help="Pretrained LoRA adapter for --ref-load." + ) + parser.add_argument( + "--teacher-lora-adapter-path", + type=str, + default=None, + help="Pretrained LoRA adapter for --teacher-load.", + ) + parser.add_argument( + "--ref-cpu-offload", + action="store_true", + help="Keep reference parameter shards in pinned CPU memory between FSDP forwards.", + ) + parser.add_argument( + "--teacher-cpu-offload", + action="store_true", + help="Keep teacher parameter shards in pinned CPU memory between FSDP forwards.", + ) return parser def add_router_arguments(parser): @@ -1577,6 +1602,25 @@ def set_default_diffusion_args(args) -> None: args.ref_mode = "none" +def validate_reference_model_args(args) -> None: + if args.ref_mode == "ref" and args.ref_load is None: + raise ValueError("--ref-mode ref requires --ref-load") + if args.ref_load is not None and args.ref_mode != "ref": + raise ValueError("--ref-load requires --ref-mode ref") + if args.ref_lora_adapter_path is not None and args.ref_load is None: + raise ValueError("--ref-lora-adapter-path requires --ref-load") + if args.teacher_lora_adapter_path is not None and args.teacher_load is None: + raise ValueError("--teacher-lora-adapter-path requires --teacher-load") + if args.ref_cpu_offload and args.ref_load is None: + raise ValueError("--ref-cpu-offload requires --ref-load") + if args.teacher_cpu_offload and args.teacher_load is None: + raise ValueError("--teacher-cpu-offload requires --teacher-load") + if args.ref_mode == "lora_base" and not args.use_lora: + raise ValueError("--ref-mode lora_base requires --use-lora") + if args.ref_mode == "ema" and not args.use_ema: + raise ValueError("--ref-mode ema requires --use-ema") + + def validate_actor_lora_adapter(args) -> None: if args.lora_adapter_path is None: return @@ -1613,6 +1657,7 @@ def validate_actor_lora_adapter(args) -> None: def miles_validate_args(args): + validate_reference_model_args(args) args.eval_datasets = _resolve_eval_datasets(args) if args.eval_interval is not None: @@ -1792,13 +1837,9 @@ def miles_validate_args(args): ) if is_nft and args.ref_mode == "none": - raise ValueError("--loss-type nft requires a reference model; set --ref-mode ema or lora_base") - if args.ref_mode == "ema" and not args.use_ema: - raise ValueError("--ref-mode ema requires --use-ema") - if args.ref_mode == "lora_base" and not args.use_lora: - raise ValueError("--ref-mode lora_base requires --use-lora") + raise ValueError("--loss-type nft requires a reference model; set --ref-mode ema, lora_base, or ref") if args.diffusion_kl_beta > 0 and args.ref_mode == "none": - raise ValueError("--diffusion-kl-beta > 0 requires a reference model; set --ref-mode lora_base or ema") + raise ValueError("--diffusion-kl-beta > 0 requires a reference model; set --ref-mode lora_base, ema, or ref") if args.dump_details is not None: args.save_debug_rollout_data = f"{args.dump_details}/rollout_data/{{rollout_id}}.pt" diff --git a/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py b/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py new file mode 100644 index 000000000..7c4166a43 --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py @@ -0,0 +1,151 @@ +"""Distributed counterpart to the independent-reference CPU contract. + + real PEFT checkpoint --> rank 0 CPU / rank 1 meta --> same FSDP mesh + | | | + +--> dense oracle +--> full / LoRA ref +--> LoRA actor + frozen gradients + AdamW + role checkpoint_path --> base component --> optional pretrained adapter + +Native CPU offload and GPU-resident references must produce the same outputs. +LoRA-base evaluation must leave adapter execution and adapter gradients intact. +""" + +import itertools +import os +import sys +from pathlib import Path +from types import SimpleNamespace + +import torch +import torch.distributed as dist +from peft import PeftModel +from tests.fast.backends.fsdp_utils.test_reference_models import ( + TinyBackend, + TinyPipelineConfig, + load_actor_forward_method, + make_actor_forward_harness, + make_model_loader_args, + save_component_lora_adapters, +) +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor + +from miles.backends.fsdp_utils.model_loader import load_fsdp_models + + +def load_dense_model(checkpoint_path, component, adapter_path, *, trainable=False): + model = TinyBackend().load_component( + component, checkpoint_path=checkpoint_path, master_dtype=torch.float32, materialize_weights=True + ) + if adapter_path is not None: + model = PeftModel.from_pretrained(model, adapter_path, subfolder=component, is_trainable=trainable) + model.train(trainable) + return model.cuda() + + +def local_tensor(tensor): + return tensor.to_local() if isinstance(tensor, DTensor) else tensor + + +def check_roles(mesh, adapter_path, reference_has_adapter, cpu_offload): + args, backend, config = make_model_loader_args(), TinyBackend(), TinyPipelineConfig() + parallel = SimpleNamespace(get_mesh=lambda name: mesh, get_optional_mesh=lambda name: None) + actor = load_fsdp_models( + args, backend, config, parallel, checkpoint_path="actor", lora_adapter_path=adapter_path, trainable=True + ) + reference = load_fsdp_models( + args, + backend, + config, + parallel, + checkpoint_path="reference", + lora_adapter_path=adapter_path if reference_has_adapter else None, + cpu_offload=cpu_offload, + ) + teacher = load_fsdp_models( + args, + backend, + config, + parallel, + checkpoint_path="teacher", + lora_adapter_path=adapter_path, + cpu_offload=cpu_offload, + ) + inputs = torch.arange(6, dtype=torch.float32, device="cuda").reshape(2, 3) / 10 + for component in args.update_weight_target_modules: + oracle = load_dense_model("actor", component, adapter_path, trainable=True) + reference_oracle = load_dense_model("reference", component, adapter_path if reference_has_adapter else None) + teacher_oracle = load_dense_model("teacher", component, adapter_path) + optimizer = torch.optim.AdamW(actor[component].parameters(), lr=0.01) + oracle_optimizer = torch.optim.AdamW(oracle.parameters(), lr=0.01) + expected_device = "cpu" if cpu_offload else "cuda" + frozen_parameters_before_update = {} + for role, model in (("ref", reference[component]), ("teacher", teacher[component])): + assert not model.training + for name, parameter in model.named_parameters(): + assert not parameter.requires_grad and local_tensor(parameter).device.type == expected_device + if cpu_offload: + assert local_tensor(parameter).is_pinned() + frozen_parameters_before_update[role, name] = local_tensor(parameter).detach().clone() + with torch.no_grad(): + target = reference_oracle(inputs) + torch.testing.assert_close(teacher[component](inputs), teacher_oracle(inputs)) + harness, observed = make_actor_forward_harness(actor, reference, component, inputs, "ref") + load_actor_forward_method()(harness, None, [], metrics=None).backward() + expected_actor_output = oracle(inputs) + (expected_actor_output - target).square().mean().backward() + torch.testing.assert_close(observed["ref_pred"], target) + torch.testing.assert_close(observed["new_pred"], expected_actor_output) + for actual, expected in zip(actor[component].parameters(), oracle.parameters(), strict=True): + if actual.requires_grad: + torch.testing.assert_close(actual.grad.full_tensor(), expected.grad, atol=2e-6, rtol=2e-6) + optimizer.step() + oracle_optimizer.step() + for actual, expected in zip(actor[component].parameters(), oracle.parameters(), strict=True): + torch.testing.assert_close(actual.full_tensor(), expected, atol=2e-6, rtol=2e-6) + for role, model in (("ref", reference[component]), ("teacher", teacher[component])): + for name, parameter in model.named_parameters(): + assert parameter.grad is None + torch.testing.assert_close(local_tensor(parameter), frozen_parameters_before_update[role, name]) + + actor[component].zero_grad(set_to_none=True) + oracle.zero_grad(set_to_none=True) + base_oracle = load_dense_model("actor", component, None) + with torch.no_grad(): + base_target = base_oracle(inputs) + forward_count = actor[component].base_model.model.forward_count.item() + harness, observed = make_actor_forward_harness(actor, {}, component, inputs, "lora_base") + load_actor_forward_method()(harness, None, [], metrics=None).backward() + expected_actor_output = oracle(inputs) + (expected_actor_output - base_target).square().mean().backward() + torch.testing.assert_close(observed["ref_pred"], base_target) + assert actor[component].training + assert not actor[component].base_model.model.block.proj.disable_adapters + assert actor[component].base_model.model.forward_count.item() == forward_count + 2 # reference + actor + for actual, expected in zip(actor[component].parameters(), oracle.parameters(), strict=True): + if actual.requires_grad: + torch.testing.assert_close(actual.grad.full_tensor(), expected.grad, atol=2e-6, rtol=2e-6) + optimizer.step() + oracle_optimizer.step() + for actual, expected in zip(actor[component].parameters(), oracle.parameters(), strict=True): + torch.testing.assert_close(actual.full_tensor(), expected, atol=2e-6, rtol=2e-6) + + +def main(): + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl", device_id=torch.device("cuda", torch.cuda.current_device())) + mesh = init_device_mesh("cuda", (dist.get_world_size(),)) + adapter_path = Path(sys.argv[1]) + if dist.get_rank() == 0: + save_component_lora_adapters(adapter_path) + dist.barrier() + for reference_has_adapter, cpu_offload in itertools.product((False, True), repeat=2): + check_roles(mesh, str(adapter_path), reference_has_adapter, cpu_offload) + dist.barrier() + if dist.get_rank() == 0: + print("OK") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/fast-gpu/backends/fsdp_utils/test_reference_models.py b/tests/fast-gpu/backends/fsdp_utils/test_reference_models.py new file mode 100644 index 000000000..6b076af91 --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/test_reference_models.py @@ -0,0 +1,43 @@ +"""Two-rank FSDP role loading and reference forwards match dense models. + + actor: saved LoRA --> FSDP --> forward/backward/AdamW --> dense gradient/weight oracle + ref: full / LoRA --> FSDP + GPU / CPUOffloadPolicy --> frozen reference output + teacher: LoRA --> FSDP + GPU / CPUOffloadPolicy --> frozen teacher output + LoRA base --> adapter-disabled reference --> unchanged actor buffers and flags + +Every role has two components; rank 0 loads checkpoints and rank 1 starts on meta. +""" + +from tests.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, suite="stage-b-5-gpu-h200", labels=["fsdp"]) + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two CUDA devices") +def test_independent_fsdp_references(tmp_path): + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + "--nproc_per_node=2", + str(Path(__file__).with_name("_reference_models_worker.py")), + str(tmp_path), + ], + env={**os.environ, "PYTHONUNBUFFERED": "1"}, + capture_output=True, + text=True, + timeout=300, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "OK" in result.stdout diff --git a/tests/fast/backends/fsdp_utils/test_reference_models.py b/tests/fast/backends/fsdp_utils/test_reference_models.py index 8299a2620..b6f71f6b2 100644 --- a/tests/fast/backends/fsdp_utils/test_reference_models.py +++ b/tests/fast/backends/fsdp_utils/test_reference_models.py @@ -1,10 +1,20 @@ -"""Loaded actor adapters preserve base-reference training and publication metadata. - - checkpoint --> component LoRA weights --> base reference --> actor gradients - loaded r/alpha -----------------------> IPC publication (ignores CLI defaults) - -Real PEFT loading and actor forward are exercised with FSDP collectives replaced. - checkpoint_path --> backend loads base; loader applies LoRA +"""Independent model roles keep actor parameters and gradients isolated. + + full / new / saved LoRA -> actor (trainable) ---> loss ---> gradients + optimizer step + +-> reference (full or LoRA) --no_grad--/ + +-> teacher (LoRA) -----------no_grad--/ + two components ---> select the matching role component + LoRA base reference ---> disable adapter ---> adapters and their gradients back on the actor + FSDP2 kept-gathered ---> lora_base forward ---> next step still trains the adapters + FSDP2 mixed precision ---> lazy init at load ---> fp32 adapter gradients whatever runs first + loaded adapter r/alpha ---> IPC publication metadata (ignores CLI init defaults) + FSDP boundary ---> TinyBlock containing the complete LoRA projection + checkpoint_path ---> backend loads base; loader applies the role's LoRA + +The CPU tests exercise real PEFT checkpoint loading and the production actor forward. +Loader tests replace FSDP collectives; the CUDA worker uses the real sharded loader. +The IPC test captures transport metadata without launching a rollout engine. +Dense model copies provide an independent output and gradient oracle. """ from tests.ci.ci_register import register_cpu_ci @@ -22,7 +32,7 @@ import torch import torch.distributed as dist import torch.testing._internal.distributed.fake_pg # noqa: F401 -from peft import LoraConfig, get_peft_model +from peft import LoraConfig, PeftModel, get_peft_model from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard @@ -171,6 +181,92 @@ def record_fsdp(model, plan, *, mesh, cpu_offload, args): return parallel, calls +@pytest.mark.parametrize("actor_mode", ["full", "new_lora", "saved_lora"]) +@pytest.mark.parametrize("reference_has_adapter", [False, True]) +def test_loaded_reference_and_teacher_are_independent_and_frozen( + tmp_path, unsharded_model_loader, actor_mode, reference_has_adapter +): + parallel, calls = unsharded_model_loader + save_component_lora_adapters(tmp_path) + args, backend, config = make_model_loader_args(), TinyBackend(), TinyPipelineConfig() + actor_has_adapter = actor_mode != "full" + args.use_lora = actor_has_adapter + args.lora_rank = 1 + args.lora_alpha = 2 + args.lora_target_modules = ["proj"] + args.lora_init_weights = "kaiming-uniform" + actor = model_loader.load_fsdp_models( + args, + backend, + config, + parallel, + checkpoint_path="actor", + lora_adapter_path=str(tmp_path) if actor_mode == "saved_lora" else None, + trainable=True, + ) + reference = model_loader.load_fsdp_models( + args, + backend, + config, + parallel, + checkpoint_path="reference", + lora_adapter_path=str(tmp_path) if reference_has_adapter else None, + cpu_offload=True, + ) + teacher = model_loader.load_fsdp_models( + args, + backend, + config, + parallel, + checkpoint_path="teacher", + lora_adapter_path=str(tmp_path), + ) + assert args.hf_checkpoint == "actor" and args.use_lora is actor_has_adapter + assert calls == [("shared-fsdp-mesh", flag) for flag in (False, False, True, True, False, False)] + assert backend.component_load_requests == [ + (component, checkpoint_path) + for checkpoint_path in ("actor", "reference", "teacher") + for component in args.update_weight_target_modules + ] + inputs = torch.arange(6, dtype=torch.float32).reshape(2, 3) / 10 + for component in args.update_weight_target_modules: + assert isinstance(actor[component], PeftModel) is actor_has_adapter + if actor_has_adapter: + adapter_config = actor[component].peft_config["default"] + assert adapter_config.r == (2 if actor_mode == "saved_lora" else args.lora_rank) + assert adapter_config.lora_alpha == (4 if actor_mode == "saved_lora" else args.lora_alpha) + assert isinstance(reference[component], PeftModel) is reference_has_adapter + oracle = copy.deepcopy(actor[component]) + frozen_parameters_before_update = { + (role, name): value.detach().clone() + for role, model in (("ref", reference[component]), ("teacher", teacher[component])) + for name, value in model.named_parameters() + } + with torch.no_grad(): + target = reference[component](inputs) + teacher_output = teacher[component](inputs) + assert not torch.equal(target, teacher_output) + harness, observed = make_actor_forward_harness(actor, reference, component, inputs, "ref") + load_actor_forward_method()(harness, None, [], metrics=None).backward() + oracle_output = oracle(inputs) + (oracle_output - target).square().mean().backward() + torch.testing.assert_close(observed["new_pred"], oracle_output) + torch.testing.assert_close(observed["ref_pred"], target) + assert not observed["ref_pred"].requires_grad + for actual, expected in zip(actor[component].parameters(), oracle.parameters(), strict=True): + if actual.requires_grad: + torch.testing.assert_close(actual.grad, expected.grad) + torch.optim.AdamW(actor[component].parameters(), lr=0.01).step() + torch.optim.AdamW(oracle.parameters(), lr=0.01).step() + for actual, expected in zip(actor[component].parameters(), oracle.parameters(), strict=True): + torch.testing.assert_close(actual, expected) + for role, model in (("ref", reference[component]), ("teacher", teacher[component])): + assert not model.training + for name, parameter in model.named_parameters(): + assert not parameter.requires_grad and parameter.grad is None + torch.testing.assert_close(parameter, frozen_parameters_before_update[role, name]) + + def test_lora_base_forward_restores_adapters_and_actor_gradients(tmp_path, unsharded_model_loader): parallel, _ = unsharded_model_loader save_component_lora_adapters(tmp_path) diff --git a/tests/fast/utils/test_reference_arguments.py b/tests/fast/utils/test_reference_arguments.py index 1bb52a3a0..3c744385f 100644 --- a/tests/fast/utils/test_reference_arguments.py +++ b/tests/fast/utils/test_reference_arguments.py @@ -1,9 +1,10 @@ -"""Pretrained actor LoRA must match checkpoint and weight-publication support. +"""Model roles are independent of the actor's full-parameter or LoRA training mode. - --use-lora --> saved adapter config --> A/B-only checkpoint + uniform IPC scaling + CLI r/alpha match - unsupported config -----------------> reject before model construction + actor: full / LoRA ---> ref: base / base + LoRA, teacher: base / base + LoRA + actor: LoRA ---> lora_base: reuse actor with its adapter disabled + actor LoRA enabled --> adapter config --> checkpoint A/B support + uniform IPC scaling + CLI r/alpha match -Validation requires the CLI rank/alpha to match the loaded adapter and leaves args unchanged. +Validation rejects incomplete role sources and unused options without changing args. """ from tests.ci.ci_register import register_cpu_ci @@ -15,7 +16,7 @@ import pytest from peft import LoraConfig -from miles.utils.arguments import validate_actor_lora_adapter +from miles.utils.arguments import validate_actor_lora_adapter, validate_reference_model_args def _args(**overrides): @@ -44,6 +45,52 @@ def _args(**overrides): return Namespace(**values) +@pytest.mark.parametrize("use_lora", [False, True]) +@pytest.mark.parametrize("role_has_adapter", [False, True]) +def test_frozen_models_do_not_inherit_actor_training_mode(use_lora, role_has_adapter): + args = _args( + use_lora=use_lora, + lora_adapter_path="actor-adapter" if use_lora else None, + ref_mode="ref", + ref_load="reference-base", + teacher_load="teacher-base", + ref_lora_adapter_path="reference-adapter" if role_has_adapter else None, + teacher_lora_adapter_path="teacher-adapter" if role_has_adapter else None, + ref_cpu_offload=True, + teacher_cpu_offload=True, + ) + original = vars(args).copy() + validate_reference_model_args(args) + assert vars(args) == original + + +def test_lora_base_needs_no_separate_reference_checkpoint(): + args = _args(use_lora=True, lora_adapter_path="actor-adapter", ref_mode="lora_base") + validate_reference_model_args(args) + assert args.ref_load is None + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"ref_mode": "ref"}, "--ref-mode ref requires --ref-load"), + ({"ref_load": "reference"}, "--ref-load requires --ref-mode ref"), + ({"ref_lora_adapter_path": "adapter"}, "--ref-lora-adapter-path requires --ref-load"), + ({"teacher_lora_adapter_path": "adapter"}, "--teacher-lora-adapter-path requires --teacher-load"), + ({"ref_cpu_offload": True}, "--ref-cpu-offload requires --ref-load"), + ({"teacher_cpu_offload": True}, "--teacher-cpu-offload requires --teacher-load"), + ({"ref_mode": "lora_base"}, "--ref-mode lora_base requires --use-lora"), + ({"ref_mode": "ema"}, "--ref-mode ema requires --use-ema"), + ], +) +def test_incomplete_model_roles_are_rejected(overrides, message): + args = _args(**overrides) + original = vars(args).copy() + with pytest.raises(ValueError, match=message): + validate_reference_model_args(args) + assert vars(args) == original + + def test_loaded_adapter_matching_actor_cli_passes_without_changing_args(tmp_path): # adapter r=2 alpha=4 == --lora-rank 2 --lora-alpha 4 -> accepted, args untouched args = _args(use_lora=True, lora_adapter_path=str(tmp_path), lora_ipc_weight_sync=True) From c07b3e8912e16407fed838656cc4d1bd44f2baf7 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 09:49:07 -0700 Subject: [PATCH 7/9] feat(fsdp): store EMA in a checkpointed optimizer --- miles/backends/fsdp_utils/actor.py | 39 ++-- miles/backends/fsdp_utils/checkpoint.py | 36 ++-- miles/backends/fsdp_utils/ema.py | 95 +++++----- miles/utils/arguments.py | 7 +- .../backends/fsdp_utils/_offload_worker.py | 98 ++++++++++ .../backends/fsdp_utils/test_offload.py | 4 + .../backends/fsdp_utils/test_ema_optimizer.py | 171 ++++++++++++++++++ .../backends/fsdp_utils/test_loss_hub_nft.py | 32 +--- 8 files changed, 372 insertions(+), 110 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/test_ema_optimizer.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 192233239..ccbb9a1a1 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -33,7 +33,7 @@ DiffusionUpdateWeightFromTensorLoRA, DiffusionUpdateWeightFromTensorLoRAIPC, ) -from .ema import EmaShadow, reshard_model +from .ema import EMAOptimizer, reshard_model from .input_dtype_policy import apply_input_dtype_policy from .loss_hub import DiffusionLossContext, flow_grpo_loss_formula, prepare_flow_grpo_batch from .lr_scheduler import get_lr_scheduler @@ -182,17 +182,16 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.global_step = 0 self.micro_step = 0 - checkpoint_payload = checkpoint.load(self) - - self.ema_shadow = None - if self.args.use_ema: - self.ema_shadow = EmaShadow( - (p for m in self.models.values() for p in m.parameters()), - decay=self.args.ema_decay_init, - uprate=self.args.ema_decay_ramp, - uphold=self.args.ema_decay_max, - flat_steps=self.args.ema_decay_flat_steps, + self.ema_optimizer = None + if args.use_ema: + self.ema_optimizer = EMAOptimizer( + self.model, + decay=args.ema_decay_init, + uprate=args.ema_decay_ramp, + uphold=args.ema_decay_max, + flat_steps=args.ema_decay_flat_steps, ) + checkpoint_payload = checkpoint.load(self) # sglang-d now supports /update_weights_from_tensor (PR #20464). if self.args.train_only: @@ -223,6 +222,8 @@ def sleep(self) -> None: for model in models.values(): offload_model(model) move_optimizer(self.optimizer, "cpu") + if self.ema_optimizer is not None: + move_optimizer(self.ema_optimizer, "cpu") clear_memory() dist.barrier(group=get_gloo_group()) print_memory("after sleep DiT") @@ -241,6 +242,8 @@ def wake_up(self) -> None: onload_model(model, cpu_offload=cpu_offload) if not self.args.fsdp_cpu_offload: move_optimizer(self.optimizer, "cuda") + if self.ema_optimizer is not None: + move_optimizer(self.ema_optimizer, "cuda") dist.barrier(group=get_gloo_group()) print_memory("after wake_up DiT") @@ -267,13 +270,10 @@ def update_weights(self) -> None: # type: ignore[override] if dist.get_rank() == 0: ray.get(self.rollout_manager.clear_num_new_engines.remote()) - ema_shadow = self.ema_shadow - if ema_shadow is not None: - delta = ema_shadow.update() - if dist.get_rank() == 0: - logger.info("EMA shadow updated (decay=%.4f step=%d)", delta, ema_shadow.step) + if self.ema_optimizer is not None: + self.ema_optimizer.step() rollout_weight_context = ( - ema_shadow.swap_in() if ema_shadow is not None and self.args.ema_rollout_policy == "ema" else nullcontext() + self.ema_optimizer.use_weights(self.model) if self.args.ema_rollout_policy == "ema" else nullcontext() ) with rollout_weight_context: self.weight_updater.update_weights() @@ -516,11 +516,8 @@ def _compute_noise_pred(model) -> torch.Tensor: with torch.no_grad(): ref_pred = _compute_noise_pred(self.reference_models[prepared.component_name]).detach() elif self.args.ref_mode == "ema": - # Drop gathered copies around the swap so neither the EMA nor the actor forward sees stale weights. - reshard_model(prepared.model) - with torch.no_grad(), self.ema_shadow.swap_in(): + with torch.no_grad(), self.ema_optimizer.use_weights(prepared.model): ref_pred = _compute_noise_pred(prepared.model).detach() - reshard_model(prepared.model) elif self.args.ref_mode == "lora_base": with torch.no_grad(), prepared.model.disable_adapter(): ref_pred = _compute_noise_pred(prepared.model).detach() diff --git a/miles/backends/fsdp_utils/checkpoint.py b/miles/backends/fsdp_utils/checkpoint.py index 4137757b3..1a7eb18fe 100644 --- a/miles/backends/fsdp_utils/checkpoint.py +++ b/miles/backends/fsdp_utils/checkpoint.py @@ -174,11 +174,7 @@ def _write_checkpoint_metadata(path: Path, metadata: dict[str, Any]) -> None: def load(actor: Any) -> dict[str, Any] | None: - """Load checkpoint from disk. - - Loads model weights and optionally optimizer state from separate directories. - This allows loading weights without optimizer or deleting optimizer before loading. - """ + """Load model and EMA state independently of the optional training optimizer state.""" load_root = actor.args.load if load_root is None: return None @@ -201,6 +197,7 @@ def load(actor: Any) -> dict[str, Any] | None: model_dir = checkpoint_dir / "model" optimizer_dir = checkpoint_dir / "optimizer" lr_scheduler_dir = checkpoint_dir / "lr_scheduler" + ema_dir = checkpoint_dir / "ema" if not model_dir.exists(): logger.info(f"[FSDP] Model checkpoint {model_dir} not found; skipping load.") @@ -218,9 +215,18 @@ def load(actor: Any) -> dict[str, Any] | None: logger.error(f"[FSDP] Failed to load model from {model_dir}: {e}") return None + if actor.ema_optimizer is not None: + if ema_dir.exists(): + ema_state = OptimizerState(actor.model, actor.ema_optimizer) + dcp.load({"ema_state": ema_state}, checkpoint_id=str(ema_dir)) + logger.info(f"[FSDP] Loaded EMA from {ema_dir}") + else: + actor.ema_optimizer.reset_from_model() + logger.info("[FSDP] EMA checkpoint missing; initialized EMA from the loaded model") + # Load optimizer state (optional) load_optimizer = not actor.args.no_load_optim - if load_optimizer and optimizer_dir.exists(): + if load_optimizer and (optimizer_dir / ".metadata").exists(): allowed_missing = actor.train_pipeline_config.optimizer_state_allowed_missing optimizer_state = OptimizerState(actor.model, actor.optimizer, allowed_missing=allowed_missing) optim_state_dict = {"optim_state": optimizer_state} @@ -233,7 +239,7 @@ def load(actor: Any) -> dict[str, Any] | None: logger.info(f"[FSDP] Optimizer checkpoint not found at {optimizer_dir}, skipping optimizer load.") # Load LR scheduler state (optional) - load_lr_scheduler = lr_scheduler_dir.exists() + load_lr_scheduler = (lr_scheduler_dir / ".metadata").exists() if load_lr_scheduler: lr_scheduler_state = LRSchedulerState(actor.lr_scheduler) lr_scheduler_state_dict = {"lr_scheduler_state": lr_scheduler_state} @@ -288,11 +294,7 @@ def finalize_load(actor: Any, checkpoint_payload: dict[str, Any] | None) -> None def save(actor: Any, iteration: int) -> None: - """Save checkpoint to disk. - - Saves model weights and optimizer state to separate directories. - This allows loading weights without optimizer or deleting optimizer before loading. - """ + """Save model and EMA state even when training optimizer state is omitted.""" torch.cuda.synchronize() base_dir = Path(actor.args.save).expanduser() @@ -301,12 +303,14 @@ def save(actor: Any, iteration: int) -> None: model_dir = checkpoint_dir / "model" optimizer_dir = checkpoint_dir / "optimizer" lr_scheduler_dir = checkpoint_dir / "lr_scheduler" + ema_dir = checkpoint_dir / "ema" if dist.get_rank() == 0: checkpoint_dir.mkdir(parents=True, exist_ok=True) model_dir.mkdir(parents=True, exist_ok=True) - optimizer_dir.mkdir(parents=True, exist_ok=True) - lr_scheduler_dir.mkdir(parents=True, exist_ok=True) + if not actor.args.no_save_optim: + optimizer_dir.mkdir(parents=True, exist_ok=True) + lr_scheduler_dir.mkdir(parents=True, exist_ok=True) dist.barrier() # Save model weights @@ -315,6 +319,10 @@ def save(actor: Any, iteration: int) -> None: state_dict = {"model_state": model_state} dcp.save(state_dict, checkpoint_id=str(model_dir)) + if actor.ema_optimizer is not None: + ema_state = OptimizerState(actor.model, actor.ema_optimizer) + dcp.save({"ema_state": ema_state}, checkpoint_id=str(ema_dir)) + # --no-save-optim drops both the optimizer and the LR scheduler. if not actor.args.no_save_optim: allowed_missing = actor.train_pipeline_config.optimizer_state_allowed_missing diff --git a/miles/backends/fsdp_utils/ema.py b/miles/backends/fsdp_utils/ema.py index aef9d0711..9832f6fc6 100644 --- a/miles/backends/fsdp_utils/ema.py +++ b/miles/backends/fsdp_utils/ema.py @@ -1,18 +1,14 @@ -"""EMA shadow of trainable parameters for diffusion FSDP training.""" - from __future__ import annotations -from collections.abc import Iterable from contextlib import contextmanager import torch -import torch.nn as nn from torch.distributed.fsdp import FSDPModule from torch.distributed.tensor import DTensor -def _local(t: torch.Tensor) -> torch.Tensor: - return t.to_local() if isinstance(t, DTensor) else t +def _local_tensor(tensor: torch.Tensor) -> torch.Tensor: + return tensor.to_local() if isinstance(tensor, DTensor) else tensor def reshard_model(model: torch.nn.Module) -> None: @@ -22,61 +18,70 @@ def reshard_model(model: torch.nn.Module) -> None: module.reshard() -class EmaShadow: - """EMA shadow of trainable parameters.""" - +class EMAOptimizer(torch.optim.Optimizer): def __init__( self, - parameters: Iterable[nn.Parameter], + model: torch.nn.Module, *, decay: float = 0.001, uprate: float = 0.001, uphold: float = 0.5, flat_steps: int = 0, ) -> None: - self.decay = float(decay) - self.uprate = float(uprate) - self.uphold = float(uphold) - self.flat_steps = int(flat_steps) - self.step = 0 - self._swapped = False + super().__init__( + (parameter for parameter in model.parameters() if parameter.requires_grad), + dict(decay=decay, uprate=uprate, uphold=uphold, flat_steps=flat_steps, update_count=0), + ) + self.reset_from_model() - self.params = [p for p in parameters if p.requires_grad] - if not self.params: - raise ValueError("EmaShadow: model has no trainable parameters") - self.shadow = [_local(p.detach()).clone() for p in self.params] + @property + def update_count(self) -> int: + return self.param_groups[0]["update_count"] - def decay_at(self, t: int) -> float: - if t <= self.flat_steps: - return self.decay - return float(min((t - self.flat_steps) * self.uprate, self.uphold)) + @torch.no_grad() + def reset_from_model(self) -> None: + for group in self.param_groups: + group["update_count"] = 0 + for parameter in group["params"]: + self.state[parameter]["ema"] = parameter.detach().clone() @torch.no_grad() - def update(self) -> float: - """theta_old <- delta * theta_old + (1 - delta) * theta.""" - if self._swapped: - raise RuntimeError("EmaShadow.update called while swapped in") - self.step += 1 - delta = self.decay_at(self.step) - for live, sh in zip(self.params, self.shadow, strict=True): - sh.mul_(delta).add_(_local(live.detach()).to(sh.device), alpha=1.0 - delta) - return delta + def step(self, closure=None) -> float: + for group in self.param_groups: + group["update_count"] += 1 + update_count = group["update_count"] + decay = ( + group["decay"] + if update_count <= group["flat_steps"] + else min((update_count - group["flat_steps"]) * group["uprate"], group["uphold"]) + ) + ema_tensors = [_local_tensor(self.state[parameter]["ema"]) for parameter in group["params"]] + actor_tensors = [_local_tensor(parameter.detach()) for parameter in group["params"]] + torch._foreach_lerp_(ema_tensors, actor_tensors, 1.0 - decay) + return decay @contextmanager - def swap_in(self): - """Temporarily expose EMA weights as the live parameters.""" - self._swap() - self._swapped = True + def use_weights(self, model: torch.nn.Module): + """Temporarily put the EMA weights into ``model``'s shards. + + Gathered full parameters are dropped before swapping in and before swapping back, + so no forward reads stale weights. + """ + reshard_model(model) + parameters = [parameter for parameter in model.parameters() if parameter in self.state] + self._swap_with_actor(parameters) try: yield finally: - self._swap() - self._swapped = False + reshard_model(model) + self._swap_with_actor(parameters) @torch.no_grad() - def _swap(self) -> None: - for live, sh in zip(self.params, self.shadow, strict=True): - live_local = _local(live.data) - tmp = live_local.clone() - live_local.copy_(sh) - sh.copy_(tmp) + def _swap_with_actor(self, parameters: list[torch.nn.Parameter]) -> None: + # EMA shares the actor's device, so its own storage holds the actor weights while they are swapped out. + for parameter in parameters: + actor_tensor = _local_tensor(parameter) + ema_tensor = _local_tensor(self.state[parameter]["ema"]) + actor_copy = actor_tensor.clone() + actor_tensor.copy_(ema_tensor) + ema_tensor.copy_(actor_copy) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7ce2b0ca9..cbd76c947 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -13,6 +13,7 @@ import argparse import json import logging +import math import os from typing import Any @@ -1763,8 +1764,8 @@ def miles_validate_args(args): if not 0.0 <= args.ema_decay_init <= 1.0: raise ValueError(f"--ema-decay-init must be in [0, 1], got {args.ema_decay_init}") - if args.ema_decay_ramp < 0.0: - raise ValueError(f"--ema-decay-ramp must be non-negative, got {args.ema_decay_ramp}") + if not math.isfinite(args.ema_decay_ramp) or args.ema_decay_ramp < 0.0: + raise ValueError(f"--ema-decay-ramp must be finite and non-negative, got {args.ema_decay_ramp}") if not 0.0 <= args.ema_decay_max <= 1.0: raise ValueError(f"--ema-decay-max must be in [0, 1], got {args.ema_decay_max}") if args.ema_decay_flat_steps < 0: @@ -1817,7 +1818,7 @@ def miles_validate_args(args): if args.ref_mode != "none": raise ValueError("--loss-type sft_loss does not use a reference model; drop --ref-mode") if args.use_ema: - raise ValueError("--loss-type sft_loss does not support --use-ema (EMA updates run in weight sync)") + raise ValueError("--loss-type sft_loss does not support --use-ema") is_nft = args.loss_type == "nft" if is_nft: diff --git a/tests/fast-gpu/backends/fsdp_utils/_offload_worker.py b/tests/fast-gpu/backends/fsdp_utils/_offload_worker.py index 106df29fd..0bf766739 100644 --- a/tests/fast-gpu/backends/fsdp_utils/_offload_worker.py +++ b/tests/fast-gpu/backends/fsdp_utils/_offload_worker.py @@ -7,16 +7,27 @@ The 2 x 2 x 2 matrix covers frozen references, cached gathered parameters, and native CPUOffloadPolicy. Uneven local shards must remain pinned after FSDP padding; optimizer Parameter bindings and tied/nonpersistent buffer aliases must survive. + + GPU/CPU actor x GPU/CPU EMA --> EMA inference --> actor forward/backward + | | + pinned sleep <--------- EMA update --------+ + | + DCP save/load --> same EMA tensors, placement, and update count """ import copy import importlib.util import itertools import os +import shutil +import sys +import tempfile from pathlib import Path import torch import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict, set_optimizer_state_dict from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import CPUOffloadPolicy, OffloadPolicy, fully_shard from torch.distributed.tensor import DTensor @@ -25,6 +36,10 @@ _SPEC = importlib.util.spec_from_file_location("fsdp_offload_under_test", _MODULE_PATH) offload = importlib.util.module_from_spec(_SPEC) _SPEC.loader.exec_module(offload) +sys.modules["miles.backends.fsdp_utils.offload"] = offload +_EMA_SPEC = importlib.util.spec_from_file_location("fsdp_ema_under_test", _MODULE_PATH.with_name("ema.py")) +ema_module = importlib.util.module_from_spec(_EMA_SPEC) +_EMA_SPEC.loader.exec_module(ema_module) class Tiny(torch.nn.Module): @@ -108,12 +123,95 @@ def check_model_offload_round_trip(mesh, frozen, reshard_after_forward, cpu_offl assert all(parameter is parameters[name] for name, parameter in model.named_parameters()) +def check_ema_offload_and_checkpoint(mesh, cpu_offload): + torch.manual_seed(31) + model = Tiny().cuda() + model.input.scale_alias = model.scale + model.input.requires_grad_(False) + control = copy.deepcopy(model) + policy = CPUOffloadPolicy() if cpu_offload else OffloadPolicy() + fully_shard(model.input, mesh=mesh, reshard_after_forward=False, offload_policy=policy) + fully_shard(model, mesh=mesh, reshard_after_forward=False, offload_policy=policy) + optimizer = torch.optim.AdamW((parameter for parameter in model.parameters() if parameter.requires_grad), lr=0.01) + control_optimizer = torch.optim.AdamW( + (parameter for parameter in control.parameters() if parameter.requires_grad), lr=0.01 + ) + ema = ema_module.EMAOptimizer(model, decay=0.5, flat_steps=10) + expected_ema = { + name: parameter.detach().clone() for name, parameter in control.named_parameters() if parameter.requires_grad + } + inputs = torch.arange(12, dtype=torch.float32, device="cuda").reshape(4, 3) / 12 + for _ in range(2): + averaged_control = copy.deepcopy(control) + with torch.no_grad(): + for name, parameter in averaged_control.named_parameters(): + if parameter.requires_grad: + parameter.copy_(expected_ema[name]) + with torch.no_grad(), ema.use_weights(model): + torch.testing.assert_close(model(inputs), averaged_control(inputs), atol=2e-6, rtol=2e-6) + + actor_output = model(inputs) + control_output = control(inputs) + torch.testing.assert_close(actor_output, control_output, atol=2e-6, rtol=2e-6) + actor_output.square().mean().backward() + control_output.square().mean().backward() + optimizer.step() + control_optimizer.step() + optimizer.zero_grad(set_to_none=True) + control_optimizer.zero_grad(set_to_none=True) + assert ema.step() == 0.5 + for name, parameter in control.named_parameters(): + if parameter.requires_grad: + expected_ema[name].lerp_(parameter.detach(), 0.5) + for name, parameter in model.named_parameters(): + if parameter.requires_grad: + average = ema.state[parameter]["ema"] + assert isinstance(average, DTensor) + expected_device = "cpu" if cpu_offload else "cuda" + assert local_tensor(average).device.type == expected_device + torch.testing.assert_close(average.to("cuda").full_tensor(), expected_ema[name], atol=2e-6, rtol=2e-6) + else: + assert parameter not in ema.state + + offload.offload_model(model) + offload.move_optimizer(optimizer, "cpu") + offload.move_optimizer(ema, "cpu") + # With native CPU offload the EMA never leaves the CPU, so only GPU-resident EMA becomes pinned. + assert all(local_tensor(state["ema"]).is_pinned() or cpu_offload for state in ema.state.values()) + offload.onload_model(model, cpu_offload=cpu_offload) + offload.move_optimizer(optimizer, "cpu" if cpu_offload else "cuda") + offload.move_optimizer(ema, "cpu" if cpu_offload else "cuda") + + checkpoint_path = [tempfile.mkdtemp(prefix="stage1_ema_") if dist.get_rank() == 0 else None] + dist.broadcast_object_list(checkpoint_path, src=0) + try: + dcp.save({"ema": get_optimizer_state_dict(model, ema)}, checkpoint_id=checkpoint_path[0]) + resumed = ema_module.EMAOptimizer(model, decay=0.1) + state_dict = {"ema": get_optimizer_state_dict(model, resumed)} + dcp.load(state_dict, checkpoint_id=checkpoint_path[0]) + set_optimizer_state_dict(model, resumed, state_dict["ema"]) + assert resumed.update_count == ema.update_count == 2 + assert resumed.param_groups[0]["decay"] == 0.5 + for parameter, state in resumed.state.items(): + average = state["ema"] + expected = ema.state[parameter]["ema"] + assert isinstance(average, DTensor) and average.placements == expected.placements + assert local_tensor(average).device == local_tensor(expected).device + torch.testing.assert_close(local_tensor(average), local_tensor(expected)) + finally: + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(checkpoint_path[0]) + + def main(): torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) dist.init_process_group("nccl", device_id=torch.device("cuda", torch.cuda.current_device())) mesh = init_device_mesh("cuda", (dist.get_world_size(),)) for frozen, reshard_after_forward, cpu_offload in itertools.product((False, True), repeat=3): check_model_offload_round_trip(mesh, frozen, reshard_after_forward, cpu_offload) + for cpu_offload in (False, True): + check_ema_offload_and_checkpoint(mesh, cpu_offload) dist.barrier() if dist.get_rank() == 0: print("OK") diff --git a/tests/fast-gpu/backends/fsdp_utils/test_offload.py b/tests/fast-gpu/backends/fsdp_utils/test_offload.py index 699aaaa3c..18bafc47f 100644 --- a/tests/fast-gpu/backends/fsdp_utils/test_offload.py +++ b/tests/fast-gpu/backends/fsdp_utils/test_offload.py @@ -5,6 +5,10 @@ uneven shards + tied/nonpersistent buffers | offload --> pinned storage --> onload --> same outputs/updates + + GPU/CPU actor + EMA on the same device --> average --> inference --> actor forward/backward + | + pinned sleep / onload + DCP round trip """ from tests.ci.ci_register import register_cuda_ci diff --git a/tests/fast/backends/fsdp_utils/test_ema_optimizer.py b/tests/fast/backends/fsdp_utils/test_ema_optimizer.py new file mode 100644 index 000000000..199915125 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_ema_optimizer.py @@ -0,0 +1,171 @@ +"""EMA is optimizer state, independent of Adam and the active model. + + dense / LoRA parameters --step--> EMA state --DCP--> resumed EMA + frozen base / buffers --------> unchanged (never averaged) + actor --use_weights--> EMA forward --finally--> original actor weights + component A --use_weights--> EMA A; component B stays untouched + +Dense and adapter-only checkpoints save EMA independently of Adam and restore its +schedule and update count. Older checkpoints seed EMA from the loaded actor weights. +""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=["fsdp"]) + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict + +from miles.backends.fsdp_utils import checkpoint +from miles.backends.fsdp_utils.ema import EMAOptimizer + + +class Model(nn.Module): + def __init__(self, *, frozen_base=False): + super().__init__() + self.base = nn.Parameter(torch.tensor([2.0]), requires_grad=not frozen_base) + self.lora_A = nn.Parameter(torch.tensor([4.0])) + self.register_buffer("counter", torch.tensor(5)) + self.register_buffer("floating_buffer", torch.tensor([8.0])) + + +@pytest.mark.parametrize("frozen_base", [False, True]) +def test_update_averages_only_trainable_parameters(frozen_base): + model = Model(frozen_base=frozen_base) + ema = EMAOptimizer(model, decay=0.5, uprate=0.1, uphold=0.9, flat_steps=1) + with torch.no_grad(): + model.base.fill_(6.0) + model.lora_A.fill_(8.0) + assert ema.step() == 0.5 + torch.testing.assert_close(ema.state[model.lora_A]["ema"], torch.tensor([6.0])) + if frozen_base: + assert model.base not in ema.state + else: + torch.testing.assert_close(ema.state[model.base]["ema"], torch.tensor([4.0])) + assert ema.step() == 0.1 + assert ema.update_count == 2 + torch.testing.assert_close(ema.state[model.lora_A]["ema"], torch.tensor([7.8])) + assert len(ema.state) == (1 if frozen_base else 2) + assert model.counter.item() == 5 + assert model.floating_buffer.item() == 8.0 + + +def test_use_weights_restores_actor_weights_after_failure(): + # actor lora_A=10, EMA lora_A=4 --use_weights--> 4 --exception--> 10 again, same Parameter object + model = Model(frozen_base=True) + original_parameter = model.lora_A + ema = EMAOptimizer(model) + with torch.no_grad(): + model.lora_A.fill_(10.0) + with pytest.raises(RuntimeError, match="reference failed"), ema.use_weights(model): + assert model.lora_A.item() == 4.0 + raise RuntimeError("reference failed") + assert model.lora_A is original_parameter and model.lora_A.item() == 10.0 + torch.testing.assert_close(ema.state[model.lora_A]["ema"], torch.tensor([4.0])) + assert model.lora_A.requires_grad and not model.base.requires_grad + + +def test_use_weights_scopes_parameters_to_selected_component(): + models = nn.ModuleDict({"transformer": Model(), "transformer_2": Model()}) + ema = EMAOptimizer(models) + with torch.no_grad(): + for parameter in models.parameters(): + parameter.add_(10.0) + first, second = models.values() + second_parameters = dict(second.named_parameters()) + second_buffers = dict(second.named_buffers()) + + with ema.use_weights(first): + assert first.base.item() == 2.0 and first.lora_A.item() == 4.0 + assert second.base.item() == 12.0 and second.lora_A.item() == 14.0 + for name, parameter in second.named_parameters(): + assert parameter is second_parameters[name] + for name, buffer in second.named_buffers(): + assert buffer is second_buffers[name] + assert first.base.item() == second.base.item() == 12.0 + assert first.lora_A.item() == second.lora_A.item() == 14.0 + + # Publishing the container still selects every component. + with ema.use_weights(models): + assert first.base.item() == second.base.item() == 2.0 + assert first.lora_A.item() == second.lora_A.item() == 4.0 + assert first.base.item() == second.base.item() == 12.0 + assert first.lora_A.item() == second.lora_A.item() == 14.0 + + +@pytest.fixture +def cpu_checkpoint(monkeypatch): + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1) + monkeypatch.setattr(torch.distributed, "barrier", lambda: None) + + +def make_actor(path, *, use_ema=True, use_lora=False): + model = nn.ModuleDict({"transformer": Model(frozen_base=use_lora), "transformer_2": Model(frozen_base=True)}) + optimizer = torch.optim.AdamW(parameter for parameter in model.parameters() if parameter.requires_grad) + return SimpleNamespace( + model=model, + optimizer=optimizer, + ema_optimizer=EMAOptimizer(model, decay=0.4, uprate=0.2, uphold=0.8, flat_steps=2) if use_ema else None, + lr_scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0), + global_step=3, + micro_step=6, + train_pipeline_config=SimpleNamespace(optimizer_state_allowed_missing=[]), + args=SimpleNamespace( + save=str(path), load=str(path), ckpt_step=None, use_lora=use_lora, no_save_optim=True, no_load_optim=True + ), + ) + + +@pytest.mark.parametrize("use_lora", [False, True]) +def test_checkpoint_restores_ema_without_adam_and_continues_schedule(tmp_path, cpu_checkpoint, use_lora): + actor = make_actor(tmp_path, use_lora=use_lora) + with torch.no_grad(): + for parameter in actor.model.parameters(): + if parameter.requires_grad: + parameter.add_(2.0) + actor.ema_optimizer.step() + actor.ema_optimizer.step() + expected = get_optimizer_state_dict(actor.model, actor.ema_optimizer) + expected_names = {"transformer.lora_A", "transformer_2.lora_A"} + if not use_lora: + expected_names.add("transformer.base") + assert set(expected["state"]) == expected_names + checkpoint.save(actor, iteration=2) + assert (tmp_path / "iter_0000003" / "ema" / ".metadata").exists() + assert not (tmp_path / "iter_0000003" / "optimizer" / ".metadata").exists() + + resumed = make_actor(tmp_path, use_lora=use_lora) + resumed.ema_optimizer.param_groups[0]["uprate"] = 0.9 + checkpoint.load(resumed) + assert resumed.ema_optimizer.update_count == 2 + assert resumed.ema_optimizer.param_groups[0]["uprate"] == 0.2 + actual = get_optimizer_state_dict(resumed.model, resumed.ema_optimizer) + for name, state in expected["state"].items(): + torch.testing.assert_close(actual["state"][name]["ema"], state["ema"]) + assert actor.ema_optimizer.step() == resumed.ema_optimizer.step() == 0.2 + actual = get_optimizer_state_dict(resumed.model, resumed.ema_optimizer) + expected = get_optimizer_state_dict(actor.model, actor.ema_optimizer) + for name, state in expected["state"].items(): + torch.testing.assert_close(actual["state"][name]["ema"], state["ema"]) + + +def test_legacy_checkpoint_seeds_ema_from_loaded_actor(tmp_path, cpu_checkpoint, caplog): + actor = make_actor(tmp_path, use_ema=False) + with torch.no_grad(): + for parameter in actor.model.parameters(): + parameter.add_(9.0) + checkpoint.save(actor, iteration=0) + resumed = make_actor(tmp_path) + resumed.ema_optimizer.step() + with caplog.at_level("INFO"): + checkpoint.load(resumed) + assert resumed.ema_optimizer.update_count == 0 + assert "initialized EMA from the loaded model" in caplog.text + for parameter, state in resumed.ema_optimizer.state.items(): + torch.testing.assert_close(state["ema"], parameter) diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py index 8fe5ea642..6798fb010 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py @@ -1,4 +1,8 @@ -"""Smoke tests for DiffusionNFT hooks (prepare + loss formula; actor owns DiT).""" +"""DiffusionNFT sampling, timestep expansion, and batch preparation. + + rollout samples --> timestep pairs --> prepared noisy latents --> NFT loss + seed + rollout + microbatch --> repeatable noise and per-sample timestep order +""" from tests.ci.ci_register import register_cpu_ci @@ -10,7 +14,6 @@ from miles.backends.fsdp_utils.configs.qwen_image import QwenImageTrainPipelineConfig from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig -from miles.backends.fsdp_utils.ema import EmaShadow from miles.backends.fsdp_utils.loss_hub.nft import corrupt, nft_r_from_advantages, prepare_nft_batch from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext from miles.ray.data_conversion_hub.nft import expand_samples_to_train_pairs, resolve_nft_sigmas @@ -261,28 +264,3 @@ def test_a_different_microbatch_draws_different_noise(self): first = prepare_nft_batch(harness._ctx(_Sd3StyleConfig()), harness._batch()) second = prepare_nft_batch(harness._ctx(_Sd3StyleConfig(), microbatch_id=1), harness._batch()) assert not torch.equal(first.latents, second.latents) - - -class TestEmaShadow: - def _model(self): - return torch.nn.Linear(4, 4, bias=False) - - def test_snapshot_and_update(self): - m = self._model() - ema = EmaShadow(m.parameters(), decay=0.5, uprate=0.001, uphold=0.5, flat_steps=10) - init = m.weight.detach().clone() - with torch.no_grad(): - m.weight.add_(1.0) - delta = ema.update() - assert delta == 0.5 - assert torch.allclose(ema.shadow[0], init + 0.5) - - def test_swap_in_restores_exactly(self): - m = self._model() - ema = EmaShadow(m.parameters(), decay=0.1) - live = m.weight.detach().clone() - with torch.no_grad(): - m.weight.add_(2.0) - with ema.swap_in(): - assert torch.equal(m.weight.detach(), live) - assert torch.equal(m.weight.detach(), live + 2.0) From 8ffacb2db7c8141757dd6f38ade8ec6329a4bdba Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 09:49:10 -0700 Subject: [PATCH 8/9] fix(fsdp): update EMA once after training instead of on every weight publication --- miles/backends/fsdp_utils/actor.py | 4 +- .../test_sd3_nft_pickscore_3xGPU.json | 34 +++++------ .../backends/fsdp_utils/test_ema_optimizer.py | 61 +++++++++++++++++++ 3 files changed, 80 insertions(+), 19 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index ccbb9a1a1..961b6ef33 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -270,8 +270,6 @@ def update_weights(self) -> None: # type: ignore[override] if dist.get_rank() == 0: ray.get(self.rollout_manager.clear_num_new_engines.remote()) - if self.ema_optimizer is not None: - self.ema_optimizer.step() rollout_weight_context = ( self.ema_optimizer.use_weights(self.model) if self.args.ema_rollout_policy == "ema" else nullcontext() ) @@ -307,6 +305,8 @@ def train(self, rollout_id: int, rollout_data_ref) -> None: # type: ignore[over if self.args.debug_rollout_only: return self._train_core(rollout_id=rollout_id, rollout_data=rollout_data) + if self.ema_optimizer is not None: + self.ema_optimizer.step() train_metric_utils.log_perf_data_raw( rollout_id=rollout_id, diff --git a/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json b/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json index b7a203791..4ae363680 100644 --- a/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json +++ b/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json @@ -1,6 +1,6 @@ { "meta": { - "commit": "29fde6eea9f11d90472d1b0d73a350a6d6739893", + "commit": "79d56ec9cffa9cb79aa9273b21d058e5c846a5d8", "source": "test_sd3_nft_pickscore_3xGPU.py" }, "metrics": { @@ -15,11 +15,11 @@ ], [ 2, - 0.7544383406639099 + 0.7545289993286133 ], [ 3, - 0.7722703218460083 + 0.7717660665512085 ] ], "rollout/reward/raw_median": [ @@ -33,11 +33,11 @@ ], [ 2, - 0.7481430172920227 + 0.7480742931365967 ], [ 3, - 0.7695608139038086 + 0.7709771990776062 ] ], "rollout/reward/raw_num_samples": [ @@ -69,11 +69,11 @@ ], [ 2, - 0.04250789433717728 + 0.04240585118532181 ], [ 3, - 0.056108858436346054 + 0.05588889122009277 ] ], "train/grad_norm": [ @@ -87,11 +87,11 @@ ], [ 3.0, - 0.028574198484420776 + 0.028748562559485435 ], [ 4.0, - 0.04176021367311478 + 0.042358506470918655 ] ], "train/nft_loss": [ @@ -105,11 +105,11 @@ ], [ 3.0, - 17.19943576388889 + 17.19867621527778 ], [ 4.0, - 18.152940538194443 + 18.16205512152778 ] ], "train/nft_neg_loss": [ @@ -123,11 +123,11 @@ ], [ 3.0, - 0.3821970621744792 + 0.38219621446397567 ], [ 4.0, - 0.4034067789713542 + 0.4036237928602431 ] ], "train/nft_pos_loss": [ @@ -141,11 +141,11 @@ ], [ 3.0, - 0.38222164577907985 + 0.38220638699001735 ], [ 4.0, - 0.4034491644965278 + 0.40365176730685765 ] ], "train/nft_r_mean": [ @@ -159,11 +159,11 @@ ], [ 3.0, - 0.49999996460974216 + 0.5000000349763367 ], [ 4.0, - 0.5000000204890966 + 0.5000000016556846 ] ], "train/nft_t_mean": [ diff --git a/tests/fast/backends/fsdp_utils/test_ema_optimizer.py b/tests/fast/backends/fsdp_utils/test_ema_optimizer.py index 199915125..b56bd427f 100644 --- a/tests/fast/backends/fsdp_utils/test_ema_optimizer.py +++ b/tests/fast/backends/fsdp_utils/test_ema_optimizer.py @@ -4,6 +4,7 @@ frozen base / buffers --------> unchanged (never averaged) actor --use_weights--> EMA forward --finally--> original actor weights component A --use_weights--> EMA A; component B stays untouched + actor.train --> one EMA update --> repeated publications (no further updates) Dense and adapter-only checkpoints save EMA independently of Adam and restore its schedule and update count. Older checkpoints seed EMA from the loaded actor weights. @@ -13,6 +14,9 @@ register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=["fsdp"]) +import ast +from contextlib import nullcontext +from pathlib import Path from types import SimpleNamespace import pytest @@ -97,6 +101,63 @@ def test_use_weights_scopes_parameters_to_selected_component(): assert first.lora_A.item() == second.lora_A.item() == 14.0 +def test_actor_train_updates_ema_once_and_publication_only_reads_it(): + path = Path(__file__).resolve().parents[4] / "miles/backends/fsdp_utils/actor.py" + actor_class = next( + node + for node in ast.parse(path.read_text()).body + if isinstance(node, ast.ClassDef) and node.name == "FSDPTrainRayActor" + ) + methods = [ + node + for node in actor_class.body + if isinstance(node, ast.FunctionDef) and node.name in ("train", "update_weights") + ] + for method in methods: + method.decorator_list = [] + namespace = dict( + ray=SimpleNamespace(get=lambda value: value), + dist=SimpleNamespace(get_rank=lambda: 0), + timer=lambda name: nullcontext(), + inverse_timer=lambda name: nullcontext(), + nullcontext=nullcontext, + clear_memory=lambda: None, + train_metric_utils=SimpleNamespace(log_perf_data_raw=lambda **kwargs: None), + ) + exec(compile(ast.Module(body=methods, type_ignores=[]), str(path), "exec"), namespace) + + model = Model() + published_weights = [] + + @torch.no_grad() + def train_core(**kwargs): + for parameter in model.parameters(): + parameter.add_(2.0) + + actor = SimpleNamespace( + args=SimpleNamespace( + offload_train=False, debug_rollout_only=False, train_only=False, ema_rollout_policy="ema" + ), + model=model, + ema_optimizer=EMAOptimizer(model, decay=0.5, flat_steps=10), + _train_core=train_core, + parallel_state=SimpleNamespace(get_mesh=lambda name: SimpleNamespace(get_local_rank=lambda: 0)), + rollout_manager=SimpleNamespace(get_rollout_engines_and_lock=SimpleNamespace(remote=lambda: (None, None, 0))), + weight_updater=SimpleNamespace(update_weights=lambda: published_weights.append(model.base.detach().clone())), + ) + namespace["train"](actor, rollout_id=0, rollout_data_ref=[SimpleNamespace(inner={})]) + assert model.base.item() == 4.0 + assert actor.ema_optimizer.update_count == 1 + assert actor.ema_optimizer.state[model.base]["ema"].item() == 3.0 + + for _ in range(2): + namespace["update_weights"](actor) + assert published_weights[-1].item() == 3.0 + assert model.base.item() == 4.0 + assert actor.ema_optimizer.update_count == 1 + assert actor.ema_optimizer.state[model.base]["ema"].item() == 3.0 + + @pytest.fixture def cpu_checkpoint(monkeypatch): monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) From 14615bfb02634b970f06c450a0a67a6f18fd610b Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 22 Sep 2026 09:50:39 -0700 Subject: [PATCH 9/9] docs(fsdp): explain model roles and offload lifecycles --- docs/fsdp_model_roles.md | 129 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 docs/fsdp_model_roles.md diff --git a/docs/fsdp_model_roles.md b/docs/fsdp_model_roles.md new file mode 100644 index 000000000..9ea12dd7c --- /dev/null +++ b/docs/fsdp_model_roles.md @@ -0,0 +1,129 @@ +# FSDP model roles and phase offload + +The FSDP trainer owns one trainable actor, optional independent frozen reference +and teacher models, and an optional EMA optimizer. Each model contains the same +configured components; Wan2.2 can therefore have both `transformer` and +`transformer_2` in every role. + +```text +train worker + actor components -------- AdamW + scheduler + | + +----------------- EMAOptimizer.state[parameter]["ema"] + reference components ---- frozen, eval, independent checkpoint + teacher components ------ frozen, eval, independent checkpoint +``` + +Only actor parameters belong to AdamW. EMA stores averages for trainable actor +parameters in a separate `torch.optim.Optimizer`; frozen base parameters and all +buffers are excluded from averaging. There is no tensor snapshot manager. + +## Model sources + +Use an independent reference with: + +```bash +--hf-checkpoint /models/student \ +--ref-mode ref --ref-load /models/reference \ +--teacher-load /models/teacher +``` + +All three sources use the same family backend, component names, input preparation, +and device mesh. This provides separate model instances, not cross-family +distillation or a new DMD2 loss. The existing loss receives the reference prediction. +Custom prepare/loss hooks can access `ctx.reference_models` and `ctx.teacher_models` +by component name. Teacher inference should run inside `torch.no_grad()`. + +Each source may load a pretrained PEFT LoRA adapter: + +```bash +--use-lora --lora-adapter-path /adapters/student --lora-rank 64 --lora-alpha 128 \ +--ref-mode ref --ref-load /models/reference \ +--ref-lora-adapter-path /adapters/reference \ +--teacher-load /models/teacher \ +--teacher-lora-adapter-path /adapters/teacher +``` + +Without a role's adapter path, that frozen role loads ordinary full weights, even +when the actor uses LoRA. Full weights and base-only weights use the same loader. +An adapter source is a PEFT checkpoint; for multiple components it contains a PEFT +checkpoint under each component subdirectory. `--lora-rank` and `--lora-alpha` must +match the actor adapter's `r` and `lora_alpha`. Actor adapter weights are loaded +before a training checkpoint is restored. Independent reference/teacher weights +are always reconstructed from their configured sources, so retain those sources +and their adapter checkpoints when resuming training. + +`--ref-mode lora_base` keeps the existing low-memory alternative: temporarily +disable the actor adapter for a gradient-free base forward. It does not construct +another base model. FSDP's lazy init runs at load time, while the adapters are +trainable, so the gradient dtype does not depend on which forward runs first. FSDP is +resharded before the adapters are re-enabled, so PEFT restores adapter gradients on +the shards rather than on gathered copies. + +## EMA and checkpointing + +`--use-ema` creates the EMA optimizer. `--ref-mode ema` evaluates its weights. +EMA updates once after a rollout's training completes, retaining the existing +DiffusionNFT update cadence. Publishing or retrying a weight update never advances +EMA. EMA evaluation and publication swap the EMA values into the actor's shards and +swap them back afterward, dropping FSDP's gathered copies on both sides so no forward +reads stale weights; the actor's backward re-gathers its restored shards. + +EMA state lives on the same device as its actor parameter. Phase sleep moves it to +pinned CPU memory together with AdamW state, and wake moves both back. + +Rollout publication pushes the current actor weights unless +`--ema-rollout-policy ema` asks for the EMA weights. + +```text +iter_0000001/ + model/ actor model state (adapter-only for LoRA) + optimizer/ AdamW state + lr_scheduler/ scheduler state + ema/ EMA optimizer state, decay schedule, and update count +``` + +EMA uses the existing distributed optimizer-state checkpoint wrapper, including +parameter names and DTensor sharding. `--no-save-optim` and `--no-load-optim` apply +to AdamW; EMA remains independently saved and restored. Loading a legacy checkpoint +without EMA seeds EMA from the loaded actor and logs that its averaging history +starts over. Disabling EMA simply omits this state from loading. + +## Two offload lifecycles + +| Setting | Parameter residence during training | Transfer boundary | +|---|---|---| +| No offload | GPU | None | +| `--offload-train` | GPU while awake | Training phase sleep/wake | +| `--fsdp-cpu-offload` | Actor shards on CPU | FSDP unshard | +| `--ref-cpu-offload` | Reference shards on CPU | FSDP unshard | +| `--teacher-cpu-offload` | Teacher shards on CPU | FSDP unshard | + +Reference and teacher CPU offload are opt-in. Phase sleep covers actor, reference, +teacher, their buffers, AdamW state, and EMA state. Sleep clears completed-step +gradients instead of transferring them. Transfers use pinned CPU +storage and enqueue asynchronous copies before synchronization; Parameter objects +and optimizer bindings are preserved. Wake respects each model's native FSDP +offload policy: native-offloaded parameter shards stay on CPU while buffers return +to GPU. + +Phase offload retains the model's own CPU tensors, without a second actor backup. +The helper lets FSDP rebuild its shard views and padding through `Module._apply`; +uneven shards remain pinned. This requires a small FSDP-internal adaptation of its +pinning flag and is covered by the registered CUDA test on supported PyTorch. + +## Focused validation + +CPU CI covers EMA averaging, DCP resume, model-source selection, loaded LoRA, +argument validation, and offload tensor identity. CUDA CI covers real two-rank +FSDP offload with uneven shards, reshard settings, frozen/trainable models, native +CPU offload, EMA residence, and checkpoint state. + +```bash +python -m pytest tests/fast/backends/fsdp_utils/test_ema_optimizer.py \ + tests/fast/backends/fsdp_utils/test_reference_models.py \ + tests/fast/backends/fsdp_utils/test_offload.py \ + tests/fast/utils/test_reference_arguments.py +python -m pytest tests/fast-gpu/backends/fsdp_utils/test_offload.py \ + tests/fast-gpu/backends/fsdp_utils/test_reference_models.py +```