Issues when using `accelerate` with `fp16`
Hi,I’m trying to useacceleratemodule to parallelize my model training. But I have troubles to use it when training models withfp16. If I load the model withtorch_dtype=torch.float16, I gotValueError: Attempting to unscale FP16 gradients.. But if I don’t load the model with half precision I will get a CUDA out of memory error. Below are the details of this problem:I’m fine tuning a 2.7B CLM, thestanford-crfm/BioMedLMmodel, on one A100 - 40GB GPU (I will be working on a much larger model, but I want to use this model to test my training process to make sure everything works as expected). I initially started with a training script withoutaccelerateand withoutTrainer. I can successfully train the model when I load the model in half precision with:# here device = 'cuda'
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device)I have to load the model in half precision otherwise I will get a CUDA out of memory error. I simplified my script and uploadedhereas a demonstration. When loading the model with half precision, it takes about 27GB GPU memory out of 40GB in the training process. It has plenty of rooms left on the GPU memory.Now I want to utilize theacceleratemodule (potentially withdeepspeedfor larger models) in my training script. I made the following changes:model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
accelerator = Accelerator(cpu=False, mixed_precision='fp16')
...
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)
...
# in training loop, I updated `lose.backward()` to:
accelerator.backward(loss)Hereis the updated script. I also configuredacceleratewithaccelerate config. Thedefault_config.yamlcan be found from the same Gist.Now when I tried to launch the script on the same machine withaccelerate launch --fp16 <script_path>. I got an errorValueError: Attempting to unscale FP16 gradients.. So I removedtorch_dtype=torch.float16from model loading and rely onaccelerateto downcast the model weight to half precision. But now I got CUDA out of memory error.To summarize:I can train the model successfully when loading it withtorch_dtype=torch.float16and not usingaccelerate.Withaccelerate, I cannot load the model withtorch_dtype=torch.float16. It givesValueError: Attempting to unscale FP16 gradients..If I don’t load the model withtorch_dtype=torch.float16and usefp16withaccelerate, I got CUDA out of memory error.So my question is: how can I train the model on a single A100 - 40GB GPU withaccelerate?
2023-06-07T11:37:12Z
Same issue, have you solved it?
2023-08-05T01:43:34Z
I encountered a similar problem and the fix for me was making the unet fp32 (and making the rest fp16)Ex)pipe = StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5',
                                                   torch_dtype=torch.float32)
    unet = pipe.components['unet']. # this is full precision
    noise_scheduler = pipe.components['scheduler']
    tokenizer = pipe.components['tokenizer']
    vae = pipe.components['vae'].requires_grad_(False)
    text_encoder = pipe.components['text_encoder'].requires_grad_(False)

    text_encoder.to(accelerator.device, dtype=torch.float16)
    vae.to(accelerator.device, dtype=torch.float16)
2023-10-10T18:27:09Z
weiqis:e (potentially withdeepspeedfor larger models) in my training script. I made the following changes:Notice in the original question, the model was loaded in as fp16 due to OOM. You started with fp32 which might not be an option.
2024-01-22T16:14:16Z
ValueError : Attemting to unscale fp16 Gradients - #13 by Kroshtan - PyTorch ForumsIf you are using amp you are not supposed to call model.half() or model.to(torch.float16) and let autocast perform the casting.This works for me