Function fsdp_wrapper

train/fsdp_utils.py:84–131

(original_model, fsdp_config, ignored_modules=[])

def fsdp_wrapper(original_model, fsdp_config, ignored_modules=[]):

    device_id = dist.get_rank() % torch.cuda.device_count()
    target_device = torch.device(f"cuda:{device_id}")  # Use same device as FSDP
    # Explicitly move ignored modules to the target device
    for module in ignored_modules:
        module.to(target_device)  

    if fsdp_config.sharding_strategy == 'HYBRID_SHARD':
        device_mesh = init_device_mesh(
            "cuda", 
            mesh_shape=(fsdp_config.num_replicate, fsdp_config.num_shard),
            mesh_dim_names=("replicate", "shard")
        )
        print("device_mesh shape:", device_mesh.mesh.shape)
    else:
        device_mesh = None
 
    
    return FSDP(
        original_model,
        auto_wrap_policy=functools.partial(
            transformer_auto_wrap_policy,
            transformer_layer_cls={
                Qwen2VLDecoderLayer,
                Qwen2VLMoEDecoderLayer,
                Qwen2VLMoTDecoderLayer,
                Dinov2WithRegistersLayer,
                Dinov2WithRegistersModel,
                Dinov2WithRegistersEmbeddings,
                DINOv3ViTLayer, DINOv3ViTEmbeddings,DINOv3ViTRopePositionEmbedding, DINOv3ViTModel,
                MLPconnector,
                PatchEmbed, VisionRotaryEmbedding, Qwen2VLVisionBlock, PatchMerger,
                Qwen2VisionTransformerPretrainedModel,
            },
        ),
        ignored_modules=ignored_modules,
        mixed_precision=MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.bfloat16, 
            buffer_dtype=torch.bfloat16,
        ),
        device_id=device_id,  
        sharding_strategy=ShardingStrategy[fsdp_config.sharding_strategy],
        backward_prefetch=BackwardPrefetch[fsdp_config.backward_prefetch],
        cpu_offload=CPUOffload(offload_params=fsdp_config.cpu_offload),
        device_mesh=device_mesh,
    )
    

class FSDPCheckpoint: