nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[perf]: add FSDP symmetric-memory training

Browse Source
This commit is contained in:
SolitaryThinker
2026-07-21 15:49:14 -07:00
parent 7f139e2b28
commit e42cfa5e5b
7 changed files with 103 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/train/configs/example.yaml
+1
View File
@@ -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 ---
fastvideo/tests/train/utils/test_config.py
+20
View File
@@ -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"}}
fastvideo/tests/train/utils/test_moduleloader_attention_backend.py
+55
View File
@@ -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"),
]
fastvideo/train/entrypoint/train.py
+6
View File
@@ -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,
fastvideo/train/utils/config.py
+6
View File
@@ -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(
fastvideo/train/utils/moduleloader.py
+14
View File
@@ -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
fastvideo/train/utils/training_config.py
+1
View File
@@ -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"