DeepSpeed Zero-3 is not compatible with low_cpu_mem_usage=True or with passing a device_map

[zhmzm]

Hi,

Thanks for sharing the code and models.

I run the following command

master_port=18765
split=forget10
model=llama2-7b
lr=2e-5
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port=$master_port forget.py --config-name=forget.yaml split=${split} batch_size=4 gradient_accumulation_steps=4 model_family=${model} lr=${lr}

Then I encounter the following issue

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
Traceback (most recent call last):
  File "/tofu/forget.py", line 145, in main
    model = AutoModelForCausalLM.from_pretrained(model_id, use_flash_attention_2=model_cfg["flash_attention2"]=="true", torch_dtype=torch.bfloat16, device_map=device_map)
  File "/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 564, in from_pretrained
    return model_class.from_pretrained(
  File "/python3.10/site-packages/transformers/modeling_utils.py", line 3195, in from_pretrained
    raise ValueError(
ValueError: DeepSpeed Zero-3 is not compatible with `low_cpu_mem_usage=True` or with passing a `device_map`.