Error running large model of esm15B using FSDP

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

Run the FSDP reproduction in examples/esm2_infer_fairscale_fsdp_cpu_offloading.py with the stated CPU offloading configuration. Read fairscale/nn/data_parallel/fully_sharded_data_parallel.py around _init_param_attributes and esm/model/esm2.py around forward to trace the invalid CUDA argument. Done means determining the configuration or code path responsible and confirming that the model call completes without the reported error.

Written by the indexing model from the issue text.