Repository: hpcaitech/ColossalAI
File: colossalai/booster/plugin/torch_fsdp_plugin.py (Python)

Code before:
                 use_orig_params: bool = False,
                 ignored_parameters: Optional[Iterable[torch.nn.Parameter]] = None,
    ):
        super().__init__()
        self.fsdp_kwargs = dict(sharding_strategy=sharding_strategy,
                               cpu_offload=cpu_offload,
                               auto_wrap_policy=auto_wrap_policy,
                               backward_prefetch=backward_prefetch,
                               mixed_precision=mixed_precision,
                               ignored_modules=ignored_modules,
                               param_init_fn=param_init_fn,
                               device_id=device_id,
                               sync_module_states=sync_module_states,
                               forward_prefetch=forward_prefetch,
                               limit_all_gathers=limit_all_gathers,
                               use_orig_params=use_orig_params,
                               ignored_parameters=ignored_parameters)

    def support_no_sync(self) -> bool:
        return True

    def control_precision(self) -> bool:
        return True

    def supported_precisions(self) -> List[str]:
        return ['fp16', 'fp16_apex', 'bf16', 'fp8']

    def control_device(self) -> bool:
        return True

    def supported_devices(self) -> List[str]:
        return ['cuda']

    def configure(
        self,
        model: nn.Module,
        optimizer: Optimizer,
        criterion: Callable = None,
        dataloader: DataLoader = None,
        lr_scheduler: LRScheduler = None,
    ) -> Tuple[Union[nn.Module, OptimizerWrapper, LRScheduler, DataLoader]]:

        model = model.cuda()
        # wrap the model with PyTorch FSDP
        model = TorchFSDPModel(model, **self.fsdp_kwargs)

        if not isinstance(optimizer, OptimizerWrapper):
            optimizer = OptimizerWrapper(optimizer)

        return model, optimizer, criterion, dataloader, lr_scheduler


Review comment (suggestion):
As FSDP controls precision itself, I think it cannot support "fp16_apex". Besides, torch doesn't support fp8 natively (see https://pytorch.org/docs/main/tensor_attributes.html?highlight=dtype#torch.dtype). I think this plugin also doesn't support "fp8".

Code after:

        def __init__(
            self,
            process_group: Optional[ProcessGroup] = None,
            sharding_strategy: Optional[ShardingStrategy] = None,
            cpu_offload: Optional[CPUOffload] = None,
            auto_wrap_policy: Optional[Callable] = None,
            backward_prefetch: Optional[BackwardPrefetch] = None,
            mixed_precision: Optional[MixedPrecision] = None,
            ignored_modules: Optional[Iterable[torch.nn.Module]] = None,
            param_init_fn: Optional[Callable[[nn.Module], None]] = None,
            device_id: Optional[Union[int, torch.device]] = None,
            sync_module_states: bool = False,
        ):
            super().__init__()
            self.fsdp_kwargs = dict(process_group=process_group,
                                    sharding_strategy=sharding_strategy,
                                    cpu_offload=cpu_offload,
                                    auto_wrap_policy=auto_wrap_policy,
                                    backward_prefetch=backward_prefetch,
                                    mixed_precision=mixed_precision,
                                    ignored_modules=ignored_modules,
                                    param_init_fn=param_init_fn,
                                    device_id=device_id,
                                    sync_module_states=sync_module_states)
    elif version.parse(torch.__version__) >= version.parse('2.0.0'):

        def __init__(
            self,
            process_group: ProcessGroupType = None,
            sharding_strategy: Optional[ShardingStrategy] = None,
            cpu_offload: Optional[CPUOffload] = None,
            auto_wrap_policy: Optional[Union[Callable, _FSDPPolicy]] = None,
            backward_prefetch: Optional[BackwardPrefetch] = BackwardPrefetch.BACKWARD_PRE,
            mixed_precision: Optional[MixedPrecision] = None,
            ignored_modules: Optional[Iterable[torch.nn.Module]] = None,
            param_init_fn: Optional[Callable[[nn.Module], None]] = None,
            device_id: Optional[Union[int, torch.device]] = None,
            sync_module_states: bool = False,
            forward_prefetch: bool = False,
            limit_all_gathers: bool = False,
            use_orig_params: bool = False,
            ignored_parameters: Optional[Iterable[torch.nn.Parameter]] = None,
        ):
            super().__init__()
            self.fsdp_kwargs = dict(process_group=process_group,
                                    sharding_strategy=sharding_strategy,
                                    cpu_offload=cpu_offload,
                                    auto_wrap_policy=auto_wrap_policy,
                                    backward_prefetch=backward_prefetch,
                                    mixed_precision=mixed_precision,