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)