🚀 A simple way to launch, train, and use PyTorch models on almost any device and distributed configuration, automatic mixed precision (including fp8), and easy-to-configure FSDP and DeepSpeed support
[ ] One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
[ ] My own task or dataset (give details below)
Reproduction
This line model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-405B-Instruct", ...) throws CUDA OOM. Any help would be appreciated!!
commandaccelerate launch --config_file zero3_config.yaml pretrain.py --num_processes=8 --multi_gpu
To be precise, I'm running this command with kubeflow, so
Traceback (most recent call last):
File "//pretrain.py", line 147, in <module>
model = AutoModelForCausalLM.from_pretrained(
File "/opt/conda/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 564, in from_pretrained
return model_class.from_pretrained(
File "/opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py", line 3916, in from_pretrained
) = cls._load_pretrained_model(
File "/opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4390, in _load_pretrained_model
new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
File "/opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py", line 938, in _load_state_dict_into_meta_model
hf_quantizer.create_quantized_param(model, param, param_name, param_device, state_dict, unexpected_keys)
File "/opt/conda/lib/python3.10/site-packages/transformers/quantizers/quantizer_bnb_4bit.py", line 217, in create_quantized_param
new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(target_device)
File "/opt/conda/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 327, in to
return self._quantize(device)
File "/opt/conda/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 291, in _quantize
w = self.data.contiguous().to(device)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.62 GiB. GPU 0 has a total capacity of 79.11 GiB of which 1.29 GiB is free. Process 194 has 0 bytes memory in use. Including non-PyTorch memory, this process has 0 bytes memory in use. Process 201 has 0 bytes memory in use. Process 195 has 0 bytes memory in use. Process 199 has 0 bytes memory in use. Process 198 has 0 bytes memory in use. Process 200 has 0 bytes memory in use. Process 196 has 0 bytes memory in use. Of the allocated memory 7.30 GiB is allocated by PyTorch, and 8.22 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (
https://pytorch.org/docs/stable/notes/cuda.html#environment-variables
)
System Info
Information
Tasks
no_trainer
script in theexamples
folder of thetransformers
repo (such asrun_no_trainer_glue.py
)Reproduction
This line
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-405B-Instruct", ...)
throws CUDA OOM. Any help would be appreciated!!pretrain.py
zero3_config.yaml
command
accelerate launch --config_file zero3_config.yaml pretrain.py --num_processes=8 --multi_gpu
To be precise, I'm running this command with kubeflow, soversions
error log
Expected behavior
Load and train Llama3.1 405B without CUDA OOM.