ViT Training precision fp16 bf16 fp32

[yipingji]

Hi,

I am wondering why timm codebase only accept mixed precision fp16 and fp32 but not bf16 and fp32?

Thanks

[rwightman]

It has also worked with mixed precsion (amp) bf16 for quite a while --amp --amp-dtype bfloat16

I did actually add support for pure bf16 and fp16 a month or two back... though I really would not recommend fp16 for vit, bf16 works well enough. Use --model-dtype bfloat16 and don't use the --amp flag.

python train.py --dataset hfds/timm/resisc45 --num-classes 45 --opt adamw --lr 1e-4 --warmup-epochs 3 --sched-on-updates  --pretrained --model fastvit_mci1 --model-dtype bfloat16
Training with a single process on 1 device (cuda).
Loading pretrained weights from Hugging Face hub (apple/mobileclip_s1_timm)
Missing keys (head.fc.weight, head.fc.bias) discovered while loading pretrained weights. This is expected if model is being adapted.
Model fastvit_mci1 created, param count:21062957
Data processing configuration for current model + dataset:
	input_size: (3, 256, 256)
	interpolation: bicubic
	mean: (0.0, 0.0, 0.0)
	std: (1.0, 1.0, 1.0)
	crop_pct: 0.95
	crop_mode: center
Created AdamW (adamw) optimizer: lr: 0.0001, betas: (0.9, 0.999), eps: 1e-08, weight_decay: 2e-05, amsgrad: False, foreach: None, maximize: False, capturable: False, differentiable: False, fused: None
AMP not enabled. Training in torch.bfloat16.
PyTorch version 2.5.0 available.
Scheduled epochs: 300 (epochs + cooldown_epochs). Warmup within epochs when warmup_prefix=False. LR stepped per update.

[rwightman]

And float32 has always worked, just don't use any amp or dtype flags.