flash-attention/training/configs/optimizer/adamw-apex-distributed.yaml

# @package train.optimizer
_target_: apex.contrib.optimizers.distributed_fused_adam.DistributedFusedAdam
adam_w_mode: True