Error running large model of esm15B using FSDP

[BKRBH]

When I use FSDP to fragment and run the large model esm2_t48_15B_UR50D on the server, my code looks like this:
import torch
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
from fairscale.nn.wrap import enable_wrap, wrap
import pandas as pd
import numpy as np
import esm

init the distributed world with world_size 1

###change pull_number every time
url = "tcp://localhost:23456"
torch.distributed.init_process_group(backend="nccl", init_method=url, world_size=1, rank=0)

download model data from the hub

#model_name = "esm2_t33_650M_UR50D"
#model_name = "esm2_t36_3B_UR50D"
model_name = "esm2_t48_15B_UR50D"
model_data, regression_data = esm.pretrained._download_model_and_regression_data(model_name)

initialize the model with FSDP wrapper

fsdp_params = dict(
mixed_precision=True,
flatten_parameters=True,
state_dict_device=torch.device("cpu"), # reduce GPU mem usage
cpu_offload=True, # enable cpu offloading
)
with enable_wrap(wrapper_cls=FSDP, **fsdp_params):
model, vocab = esm.pretrained.load_model_and_alphabet_core(
model_name, model_data, regression_data
)
batch_converter = vocab.get_batch_converter()
model.eval()

# Wrap each layer in FSDP separately
for name, child in model.named_children():
    if name == "layers":
        for layer_name, layer in child.named_children():
            wrapped_layer = wrap(layer)
            setattr(child, layer_name, wrapped_layer)
model = wrap(model)

data = [
("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
]

batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_tokens = batch_tokens.cuda()
print("batch_tokens:",batch_tokens)
with torch.no_grad():
results = model(tokens=batch_tokens,repr_layers=[48],return_contacts=True)###

token_representations = results["representations"][48]

print("result!!!!:",results)

Generate per-sequence representations via averaging

NOTE: token 0 is always a beginning-of-sequence token, so the first residue is token 1.

batch_lens=(batch_tokens != vocab.padding_idx).sum(1)
sequence_representations = []
for i, tokens_len in enumerate(batch_lens):
sequence_representations.append(token_representations[i, 1 : tokens_len - 1].mean(0))

file_path = "./examples/out_all/sequence_representations.pt"
torch.save(sequence_representations, file_path)

But the operation encountered the following error problem, I don't know if it was caused by improper use of FSDP：
Traceback (most recent call last):
File "./examples/esm2_infer_fairscale_fsdp_cpu_offloading.py", line 54, in
results = model(tokens=batch_tokens,repr_layers=[48],return_contacts=True)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 1327, in forward
outputs = self.module(*args, **kwargs)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/fairscale/nn/misc/flatten_params_wrapper.py", line 461, in forward
return self.module(*inputs, **kwinputs)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/esm/model/esm2.py", line 115, in forward
need_head_weights=need_head_weights,
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 1303, in forward
self._lazy_init()
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 1102, in _lazy_init
self._init_param_attributes(p)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/home/inspur/myenvs/esmfold/lib/python3.7/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 1212, in _init_param_attributes
p._cpu_grad = torch.zeros_like(p.data, device="cpu").pin_memory()
RuntimeError: CUDA error: invalid argument