[BUG] Getting this error: NotImplementedError: Cannot copy out of meta tensor; no data!

Description:
I am loading Gemma 2 on 8 GPU single node using deepspeed. When i load the model using the following code:

from transformers import AutoTokenizer, AutoModelForCausalLM
import deepspeed

QUANTIZATION TO FP4

torch_dtype = torch.bfloat16
quant_storage_dtype = torch.bfloat16

quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch_dtype,
bnb_4bit_quant_storage=quant_storage_dtype,
)
###########################

with deepspeed.zero.Init():
model = AutoModelForCausalLM.from_pretrained(
"google/gemma-2-2b-it",
quantization_config=quantization_config,
attn_implementation="eager",
torch_dtype=quant_storage_dtype,
use_cache=False
)

I get the error:
NotImplementedError: Cannot copy out of meta tensor; no data!

From the looks of it, looks like it is trying to move a meta tensor which are abstract. How to solve this issue?

The following is my deepspeed config:
ds_config = {
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": True,
# "ratio": 0.9
},
"offload_param": {
"device": "cpu",
"pin_memory": True
},
"overlap_comm": True,
"contiguous_gradients": True,
"sub_group_size": 1e3,
"reduce_bucket_size": 1e3,
"stage3_prefetch_bucket_size": 1e3,
"stage3_param_persistence_threshold": 1e3,
"stage3_max_live_parameters": 1e3,
"stage3_max_reuse_distance": 1e3,
"stage3_gather_16bit_weights_on_model_save": True
}
}

In Training Arguments, I have the following:
training_args = TrainingArguments(
output_dir="D:/gemma-2-2b-it-training-artifacts/",
report_to="none",
learning_rate=0.0002,
lr_scheduler_type='cosine',
num_train_epochs=1,
per_device_train_batch_size=1,
per_device_eval_batch_size=1,
gradient_accumulation_steps=1,
optim='adamw_torch',
logging_steps=10,
save_strategy='epoch',
max_grad_norm=0.3,
warmup_ratio=0.1,
bf16=True,
tf32=True,
gradient_checkpointing=True,
gradient_checkpointing_kwargs={"use_reentrant": True},
deepspeed=ds_config
)

I am running the script using deepspeed launcher from command line only. Cannot share script as it is on my work laptop but the error is here when offloading parameter. Need help with this please. I am not very familiar with deep speed. Also, do i need to setup FSDP for this?

Start by reproducing the shown AutoModelForCausalLM.from_pretrained call inside deepspeed.zero.Init with the Gemma 2 quantization and ZeRO-3 offload configuration. Inspect the resulting traceback to identify where the meta tensor is moved; done means the model loads successfully under the reported eight-GPU setup without the NotImplementedError.