
```python
apply_activation_checkpointing_wrapper(
       model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn
   )

Important note - this must be run after the model has been initialized with FSDP.
However, hopefully you’ve seen how some initial tuning with FSDP options can have a large impact on your training performance. 
With that, we turn our attention from how to scale within FSDP, to how to scale your server hardware for FSDP using AWS.
Large Scale Training with FSDP on AWS - For multi-node prioritize high speed network
AWS provides several services that can be used to run distributed training with FSDP: Amazon EC2 Accelerated Computing instances, AWS ParallelCluster, and Amazon Sagemaker.