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)