Overriding DeepSpeed autocasting

[gahdritz]

What's the cleanest way to manually prevent up/downcasting for certain operations with DeepSpeed fp16/bfloat16 mixed precision enabled? I'm basically looking for something equivalent to:

with torch.cuda.amp.autocast(enabled=False):
    <do stuff without autocasting>

Ultimately, what I want to have happen is that softmaxes & layernorms are forced to run using bfloat16. If there's an even cleaner global way to do that, that would also be nice.

[gahdritz]

Bumping this---this would still be hugely useful for me.

[Zyriix]

this is important, anyone can help with this? this is a good repo and thanks a lot!

[RoyJames]

Upvoting for this option!

[mzamini92]

There are some ways I dont know if helps you or not.
prevent DeepSpeed from casting softmaxes and layernorms to fp32. something like this:

import deepspeed.zero.init

with deepspeed.zero.Init(config=ds_config):
    # Prevent softmaxes and layernorms from being casted to fp32
    operations = ["softmax", "layernorm"]
    deepspeed.zero.prevent_cast(operations)

    # Do stuff without autocasting

Or enable fp16 for softmaxes and layernorms. something like this:

import deepspeed.zero.init

with deepspeed.zero.Init(config=ds_config):
    # Enable fp16 for softmaxes and layernorms
    operations = ["softmax", "layernorm"]
    fp16.enable(operations)

    # Do stuff without autocasting

Or custom cast function can return True to prevent the cast, or False to allow the cast. something like this

def custom_cast(tensor, dtype):
    operations = ["softmax", "layernorm"]
    if tensor.type() in operations:
        return True
    else:
        return False

with deepspeed.zero.Init(config=ds_config):
    # Use the custom cast function
    fp16.enable_custom_cast(custom_cast)

    # Do stuff without autocasting