(
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,
)
| 10 | |
| 11 | |
| 12 | def shard_model( |
| 13 | model, |
| 14 | device_id, |
| 15 | param_dtype=torch.bfloat16, |
| 16 | reduce_dtype=torch.float32, |
| 17 | buffer_dtype=torch.float32, |
| 18 | process_group=None, |
| 19 | sharding_strategy=ShardingStrategy.FULL_SHARD, |
| 20 | sync_module_states=True, |
| 21 | ): |
| 22 | model = FSDP( |
| 23 | module=model, |
| 24 | process_group=process_group, |
| 25 | sharding_strategy=sharding_strategy, |
| 26 | auto_wrap_policy=partial( |
| 27 | lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks), |
| 28 | # mixed_precision=MixedPrecision( |
| 29 | # param_dtype=param_dtype, |
| 30 | # reduce_dtype=reduce_dtype, |
| 31 | # buffer_dtype=buffer_dtype), |
| 32 | device_id=device_id, |
| 33 | sync_module_states=sync_module_states) |
| 34 | return model |
| 35 | |
| 36 | |
| 37 | def free_model(model): |
nothing calls this directly
no outgoing calls
no test coverage detected