apply_fsdp() — pytorch Function Reference

Architecture documentation for the apply_fsdp() function in dist_util.py from the pytorch codebase.

Function python DynamoBenchmarks Configuration calls 1 called by 1

Entity Profile

DynamoBenchmarks→ Configuration→ apply_fsdp() — pytorch Function Reference

Dependency Diagram

graph TD
  ff4518e9_afbb_35ef_3d8f_c605ec26b945["apply_fsdp()"]
  24e4ab44_f8ed_b835_f57a_8a13ebaabe74["run_model()"]
  24e4ab44_f8ed_b835_f57a_8a13ebaabe74 -->|calls| ff4518e9_afbb_35ef_3d8f_c605ec26b945
  aa6a976a_820e_3084_1e1e_2603805389df["fsdp_checkpointing_base()"]
  ff4518e9_afbb_35ef_3d8f_c605ec26b945 -->|calls| aa6a976a_820e_3084_1e1e_2603805389df
  style ff4518e9_afbb_35ef_3d8f_c605ec26b945 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/dist_util.py lines 131–142

def apply_fsdp(args, model, use_checkpointing=False, use_wrap_policy=True):
    wrap_policy = None
    blocks = MODEL_FSDP_WRAP[
        "toy_model" if model.__class__ is ToyModel else args.torchbench_model
    ]
    if use_wrap_policy:
        wrap_policy = ModuleWrapPolicy(blocks)

    model = FSDP(model, auto_wrap_policy=wrap_policy, use_orig_params=True)
    if use_checkpointing:
        fsdp_checkpointing_base(model, blocks)
    return model