| |
| from functools import partial |
|
|
| import torch |
| from torch.distributed.fsdp import FullyShardedDataParallel as FSDP |
| from torch.distributed.fsdp import MixedPrecision, ShardingStrategy |
| from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy |
|
|
|
|
| def shard_model( |
| model, |
| device_id, |
| param_dtype=torch.bfloat16, |
| reduce_dtype=torch.float32, |
| buffer_dtype=torch.float32, |
| process_group=None, |
| sharding_strategy=ShardingStrategy.FULL_SHARD, |
| sync_module_states=True, |
| ): |
| model = FSDP( |
| module=model, |
| process_group=process_group, |
| sharding_strategy=sharding_strategy, |
| auto_wrap_policy=partial( |
| lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks), |
| mixed_precision=MixedPrecision( |
| param_dtype=param_dtype, |
| reduce_dtype=reduce_dtype, |
| buffer_dtype=buffer_dtype), |
| device_id=device_id, |
| sync_module_states=sync_module_states) |
| return model |
|
|
