(
model: nn.Module,
device_mesh: DeviceMesh,
global_device_mesh: DeviceMesh,
dtype: torch.dtype,
grad_dtype: torch.dtype = None,
strategy: str = 'op_grad',
)
| 48 | |
| 49 | |
| 50 | def fsdp_transformer_ulysses( |
| 51 | model: nn.Module, |
| 52 | device_mesh: DeviceMesh, |
| 53 | global_device_mesh: DeviceMesh, |
| 54 | dtype: torch.dtype, |
| 55 | grad_dtype: torch.dtype = None, |
| 56 | strategy: str = 'op_grad', |
| 57 | ) -> FSDP: |
| 58 | tp_mesh = device_mesh['tp'] |
| 59 | if tp_mesh.size() > 1: |
| 60 | model = model.__class__.sequence_parallelize_ulysses(model, tp_mesh) |
| 61 | |
| 62 | fn = lambda m: m in model.get_fsdp_wrap_module_list() |
| 63 | strategy = strategy_mapping[strategy] |
| 64 | model = FSDP( |
| 65 | model, |
| 66 | auto_wrap_policy=partial( |
| 67 | lambda_auto_wrap_policy, |
| 68 | lambda_fn=fn, |
| 69 | ), |
| 70 | device_mesh=global_device_mesh, |
| 71 | sharding_strategy=strategy, |
| 72 | backward_prefetch=BackwardPrefetch.BACKWARD_PRE, |
| 73 | mixed_precision=MixedPrecision( |
| 74 | param_dtype=dtype, |
| 75 | reduce_dtype=grad_dtype or dtype, |
| 76 | buffer_dtype=dtype, |
| 77 | ), |
| 78 | device_id=torch.cuda.current_device(), |
| 79 | sync_module_states=True, |
| 80 | limit_all_gathers=True, |
| 81 | use_orig_params=True, |
| 82 | ) |
| 83 | # non_reentrant_wrapper = partial( |
| 84 | # checkpoint_wrapper, |
| 85 | # checkpoint_impl=CheckpointImpl.NO_REENTRANT, |
| 86 | # ) |
| 87 | # apply_activation_checkpointing( |
| 88 | # model, |
| 89 | # checkpoint_wrapper_fn=non_reentrant_wrapper, |
| 90 | # check_fn=lambda m: m in model.get_fsdp_wrap_module_list() |
| 91 | # ) |
| 92 | # model.enable_gradient_checkpointing() |
| 93 | torch.cuda.synchronize() |
| 94 | return model |
nothing calls this directly
no test coverage detected