Skip to content

Best practices for weight dtype and autocasting #438

Closed
Closed
Best practices for weight dtype and autocasting#438

Description

@c0g
c0g
opened on Sep 19, 2023

I might be confused here, but I'm curious of the following:

It seems like if we use TE's autocast with model dtype float32, we don't get to use FlashAttention since it requires the dtype to be bfloat16. If we set the model dtype to bfloat16 TE casting to FP8 will work, and we get FlashAttention, but then according to Gopher we might have convergence issues.

Is this a real concern, and is there a recipe to fix this? Can TE's autocast also support casting to BFloat16 when possible?

Activity

  1. c0g commented on Sep 19, 2023

    @c0g
    c0g
    on Sep 19, 2023
    Author

    MosaicML's blog post says to do both!

    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    # Run training step
  2. c0g
    closed this as completedon Sep 19, 2023

  3. c0g commented on Sep 27, 2023

    @c0g
    c0g
    on Sep 27, 2023
    Author

    Hmm actually found this comment: #401 (comment), which cast doubt on this conclusion.

    Does TE work with autocast?

  4. c0g
    reopened this on Sep 27, 2023

  5. timmoon10 commented on Oct 2, 2023

    @timmoon10
    timmoon10
    on Oct 2, 2023
    Member

    Sorry for the confusion, I'm confused myself. If I remember correctly, I made that comment since I ran into issues when combining te.fp8_autocast and torch.autocast, but I haven't been able to reproduce those issues. I suppose combining multiple autocast types might lead to multiple casts (fp32 -> bf16 -> fp8), but Flash Attention/cuDNN attention may be worth that overhead.

  6. ptrendx commented on Oct 12, 2023

    @ptrendx
    ptrendx
    on Oct 12, 2023
    Member

    Hi @c0g, does @timmoon10 's comment clear the confusion? TE works with autocast - please open an issue if you see any problem with it :-).

  7. ptrendx
    closed this as completedon Oct 12, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions