(model: nn.Module, device_mesh: DeviceMesh,
dtype: torch.dtype)
| 21 | |
| 22 | |
| 23 | def fsdp_text_encoder(model: nn.Module, device_mesh: DeviceMesh, |
| 24 | dtype: torch.dtype) -> FSDP: |
| 25 | if hasattr(model, 'encoder'): |
| 26 | fn = lambda m: m in list(model.encoder.block) |
| 27 | elif hasattr(model, 'text_model'): |
| 28 | fn = lambda m: m in list(model.text_model.encoder.layers) |
| 29 | else: |
| 30 | raise ValueError('unknown text encoder') |
| 31 | |
| 32 | model = FSDP( |
| 33 | model, |
| 34 | auto_wrap_policy=partial( |
| 35 | lambda_auto_wrap_policy, |
| 36 | lambda_fn=fn, |
| 37 | ), |
| 38 | device_mesh=device_mesh, |
| 39 | sharding_strategy=ShardingStrategy.HYBRID_SHARD, |
| 40 | mixed_precision=MixedPrecision(param_dtype=dtype, ), |
| 41 | device_id=torch.cuda.current_device(), |
| 42 | sync_module_states=True, |
| 43 | limit_all_gathers=True, |
| 44 | use_orig_params=True, |
| 45 | ) |
| 46 | torch.cuda.synchronize() |
| 47 | return model |
| 48 | |
| 49 | |
| 50 | def fsdp_transformer_ulysses( |
nothing calls this directly
no outgoing calls
no test coverage detected