Skip to content

FusedAdam doesn't allocate master weights for bfloat16 #1728

Open
@cbcase

Description

@cbcase

If you look at https://github.com/NVIDIA/apex/blob/master/apex/optimizers/fused_adam.py#L184, the FusedAdam code doesn't allocate master weights when the parameter dtype is bfloat16, even if you set master_weights=True (and subsequently no master weights are passed to the kernel at https://github.com/NVIDIA/apex/blob/master/apex/optimizers/fused_adam.py#L235).

Is there a specific reason for this, or is it simply an oversight?

(Seeing as bfloat has fewer mantissa bits than fp16, fp32 master weights are even more important for bfloat16 -- though, really, they are a necessity for both.)

Metadata

Metadata

Assignees

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