precompute_float8_dynamic_scale_for_fsdp

torchao.float8.precompute_float8_dynamic_scale_for_fsdp(module: Module) → None

动态计算所有 float8 参数的 scale。应在优化器步骤之后运行。它执行一次 all-reduce 来计算所有 float8 权重的 scale。示例用法

model(input).sum().backward() optim.step() precompute_float8_dynamic_scale_for_fsdp(model)