verl/utils/fsdp_utils.py

• init_fn · function · L29-L33 — def init_fn(x: torch.nn.Module)
• get_init_weight_context_manager · function · L36-L43 — def get_init_weight_context_manager(use_meta_tensor=True)
• get_fsdp_wrap_policy · function · L48-L105 — def get_fsdp_wrap_policy(module, config=None, is_lora=False)
• lambda_policy_fn · function · L75-L79 — def lambda_policy_fn(module)
• offload_fsdp_grad · function · L108-L112 — def offload_fsdp_grad(module)
• load_fsdp_grad · function · L115-L119 — def load_fsdp_grad(module, device_id)
• offload_fsdp_param_and_grad · function · L122-L129 — def offload_fsdp_param_and_grad(module, offload_grad=False)
• load_fsdp_param_and_grad · function · L132-L139 — def load_fsdp_param_and_grad(module, device_id, load_grad=False)
• offload_fsdp_optimizer · function · L142-L149 — def offload_fsdp_optimizer(optimizer)
• load_fsdp_optimizer · function · L152-L159 — def load_fsdp_optimizer(optimizer, device_id)
• meta_device_init · function · L163-L191 — def meta_device_init()
• register_empty_parameter · function · L175-L184 — def register_empty_parameter(module, name, param)
• parallel_load_safetensors · function · L194-L247 — def parallel_load_safetensors(filepath)
• parallel_init_module_fn · function · L250-L329 — def parallel_init_module_fn(module: torch.nn.Module, shard_states: Dict[str, torch.nn.Parameter])
• create_and_sync_state · function · L272-L289 — def create_and_sync_state(param_name, state, is_param)
• init_fn · function · L291-L327 — def init_fn(sub_mod: torch.nn.Module, recurse: bool = True)