From e42cfa5e5bd7223af2a7fe1632fb23478bcfc504 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 21 Jul 2026 15:49:14 -0700 Subject: [PATCH] [perf]: add FSDP symmetric-memory training --- examples/train/configs/example.yaml | 1 + fastvideo/tests/train/utils/test_config.py | 20 +++++++ .../test_moduleloader_attention_backend.py | 55 +++++++++++++++++++ fastvideo/train/entrypoint/train.py | 6 ++ fastvideo/train/utils/config.py | 6 ++ fastvideo/train/utils/moduleloader.py | 14 +++++ fastvideo/train/utils/training_config.py | 1 + 7 files changed, 103 insertions(+) diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index d80255acd..ee8150751 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -87,6 +87,7 @@ training: hsdp_shard_dim: 8 # default: -1 (defaults to num_gpus in loader) pin_cpu_memory: false # default: false reshard_after_forward: true # default: true; false trades memory for fewer FSDP all-gathers + fsdp_symmetric_memory: false # default: false; native NCCL symmetric-memory FSDP collectives reduce_dtype: fp32 # default: "fp32"; "bf16" reduces FSDP communication and memory # --- training.data [TYPED] -> DataConfig --- diff --git a/fastvideo/tests/train/utils/test_config.py b/fastvideo/tests/train/utils/test_config.py index 952e23b50..2f9ca6310 100644 --- a/fastvideo/tests/train/utils/test_config.py +++ b/fastvideo/tests/train/utils/test_config.py @@ -62,6 +62,7 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None: assert t.distributed.hsdp_replicate_dim == 1 assert t.distributed.pin_cpu_memory is False assert t.distributed.reshard_after_forward is True + assert t.distributed.fsdp_symmetric_memory is False assert t.distributed.reduce_dtype == "fp32" assert t.data.train_batch_size == 1 @@ -107,6 +108,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: "hsdp_shard_dim": 2, "pin_cpu_memory": True, "reshard_after_forward": False, + "fsdp_symmetric_memory": True, "reduce_dtype": "bf16", }, "data": { @@ -162,6 +164,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: assert t.distributed.tp_size == 2 assert t.distributed.pin_cpu_memory is True assert t.distributed.reshard_after_forward is False + assert t.distributed.fsdp_symmetric_memory is True assert t.distributed.reduce_dtype == "bf16" assert t.data.train_batch_size == 2 @@ -307,6 +310,7 @@ def test_dotted_overrides_apply_with_type_coercion(tmp_path: Path) -> None: "--training.optimizer.fused=true", "--training.distributed.pin_cpu_memory=true", "--training.distributed.reshard_after_forward=false", + "--training.distributed.fsdp_symmetric_memory=true", "--training.distributed.reduce_dtype=bf16", "--training.model.enable_torch_compile=true", "--training.tracker.project_name=overridden", @@ -318,6 +322,7 @@ def test_dotted_overrides_apply_with_type_coercion(tmp_path: Path) -> None: assert cfg.training.optimizer.fused is True assert cfg.training.distributed.pin_cpu_memory is True assert cfg.training.distributed.reshard_after_forward is False + assert cfg.training.distributed.fsdp_symmetric_memory is True assert cfg.training.distributed.reduce_dtype == "bf16" assert cfg.training.model.enable_torch_compile is True assert cfg.training.tracker.project_name == "overridden" @@ -355,6 +360,21 @@ def test_reshard_after_forward_rejects_non_bool(tmp_path: Path) -> None: load_run_config(_write_yaml(tmp_path, data)) +def test_fsdp_symmetric_memory_rejects_non_bool(tmp_path: Path) -> None: + data = _minimal_yaml() + data["training"] = { + "distributed": { + "fsdp_symmetric_memory": "true" + } + } + + with pytest.raises( + ValueError, + match="training.distributed.fsdp_symmetric_memory must be a bool", + ): + load_run_config(_write_yaml(tmp_path, data)) + + def test_model_torch_compile_rejects_non_bool(tmp_path: Path) -> None: data = _minimal_yaml() data["training"] = {"model": {"enable_torch_compile": "true"}} diff --git a/fastvideo/tests/train/utils/test_moduleloader_attention_backend.py b/fastvideo/tests/train/utils/test_moduleloader_attention_backend.py index 0815a55e3..5ca39ab4e 100644 --- a/fastvideo/tests/train/utils/test_moduleloader_attention_backend.py +++ b/fastvideo/tests/train/utils/test_moduleloader_attention_backend.py @@ -129,3 +129,58 @@ def test_load_transformer_configures_recursive_fsdp_reshard( ) assert calls == [(False, True)] + + +def test_load_transformer_configures_all_fsdp_symmetric_memory_modules( + monkeypatch, + tmp_path, +) -> None: + training_config = TrainingConfig( + distributed=DistributedConfig( + hsdp_shard_dim=1, + fsdp_symmetric_memory=True, + ), + pipeline_config=PipelineConfig(), + ) + calls: list[tuple[str, str, object]] = [] + + class FakeFSDPModule(torch.nn.Module): + + def __init__(self, name: str, child: torch.nn.Module | None = None) -> None: + super().__init__() + self.name = name + if child is not None: + self.child = child + + def set_force_sum_reduction_for_comms(self, value: bool) -> None: + calls.append((self.name, "force_sum", value)) + + def set_symm_mem_for_comm(self, backend: str) -> None: + calls.append((self.name, "symmetric_memory", backend)) + + transformer = FakeFSDPModule("root", FakeFSDPModule("block")) + monkeypatch.setattr(moduleloader, "FSDPModule", FakeFSDPModule) + monkeypatch.setattr(moduleloader, "maybe_download_model", lambda path: str(tmp_path)) + monkeypatch.setattr( + moduleloader, + "verify_model_config_and_directory", + lambda path: {"transformer": ("diffusers", "FakeTransformer")}, + ) + monkeypatch.setattr( + moduleloader.PipelineComponentLoader, + "load_module", + lambda **kwargs: transformer, + ) + + moduleloader.load_module_from_path( + model_path="fake/model", + module_type="transformer", + training_config=training_config, + ) + + assert calls == [ + ("root", "force_sum", True), + ("root", "symmetric_memory", "NCCL"), + ("block", "force_sum", True), + ("block", "symmetric_memory", "NCCL"), + ] diff --git a/fastvideo/train/entrypoint/train.py b/fastvideo/train/entrypoint/train.py index 3f71e2a72..28ccf648e 100644 --- a/fastvideo/train/entrypoint/train.py +++ b/fastvideo/train/entrypoint/train.py @@ -69,6 +69,12 @@ def run_training_from_config( "SLA_ATTN", ) + if tc.distributed.fsdp_symmetric_memory: + nccl_cta_policy = os.environ.setdefault("NCCL_CTA_POLICY", "2") + if nccl_cta_policy != "2": + raise ValueError("training.distributed.fsdp_symmetric_memory requires " + "NCCL_CTA_POLICY=2 before distributed initialization") + maybe_init_distributed_environment_and_model_parallel( tc.distributed.tp_size, tc.distributed.sp_size, diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index 2a29627ea..507e563c7 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -381,6 +381,12 @@ def _build_training_config( default=True, where="training.distributed.reshard_after_forward", ), + fsdp_symmetric_memory=require_bool( + d, + "fsdp_symmetric_memory", + default=False, + where="training.distributed.fsdp_symmetric_memory", + ), reduce_dtype=reduce_dtype, ), data=DataConfig( diff --git a/fastvideo/train/utils/moduleloader.py b/fastvideo/train/utils/moduleloader.py index ebe270eeb..b00fb5c3b 100644 --- a/fastvideo/train/utils/moduleloader.py +++ b/fastvideo/train/utils/moduleloader.py @@ -7,6 +7,7 @@ from contextlib import nullcontext from typing import Any, TYPE_CHECKING import torch +from torch.distributed.fsdp import FSDPModule from fastvideo.attention.selector import ( coerce_attn_backend, @@ -14,6 +15,7 @@ from fastvideo.attention.selector import ( ) from fastvideo.configs.pipelines.base import PipelineConfig from fastvideo.fastvideo_args import ExecutionMode, TrainingArgs +from fastvideo.logger import init_logger from fastvideo.models.loader.component_loader import ( PipelineComponentLoader, ) from fastvideo.utils import ( @@ -26,6 +28,8 @@ if TYPE_CHECKING: from fastvideo.train.utils.training_config import ( TrainingConfig, ) +logger = init_logger(__name__) + # ------------------------------------------------------------------ # TrainingArgs builders (only place that creates FastVideoArgs) # ------------------------------------------------------------------ @@ -179,4 +183,14 @@ def load_module_from_path( elif not reshard_after_forward: raise RuntimeError("training.distributed.reshard_after_forward=false requires " "an FSDP-wrapped transformer") + + if training_config.distributed.fsdp_symmetric_memory: + fsdp_modules = [submodule for submodule in module.modules() if isinstance(submodule, FSDPModule)] + if not fsdp_modules: + raise RuntimeError("training.distributed.fsdp_symmetric_memory=true requires " + "an FSDP-wrapped transformer") + for fsdp_module in fsdp_modules: + fsdp_module.set_force_sum_reduction_for_comms(True) + fsdp_module.set_symm_mem_for_comm("NCCL") + logger.info("Enabled FSDP symmetric-memory communication on %d modules", len(fsdp_modules)) return module diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 1d7489547..a7633418f 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -19,6 +19,7 @@ class DistributedConfig: hsdp_shard_dim: int = -1 pin_cpu_memory: bool = False reshard_after_forward: bool = True + fsdp_symmetric_memory: bool = False reduce_dtype: str = "fp32"