Can I use bfloat16 when training?

[yangyuya]

I see the code use torch.float by default.

    model = VideoChatGPTLlamaForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        # torch_dtype=torch.bfloat16 if training_args.bf16 else torch.float,
    )

Can I use bfloat16 when training? I find if I use bfloat16 I can train with 24G GPUs. But I'm not sure how much this affects model performance? Can you give me some advice?

[mmaaz60]

Hi @yangyuya,

Thank you for you interest in our work. I have not tried it though the training in bf16 mode should work and give similar results.