torchao.float8

主要的 float8 训练 API

convert_to_float8_training

将 module 中的 torch.nn.Linear 替换为 Float8Linear。

其他 float8 训练类型

Float8LinearConfig	用于将 torch.nn.Linear 模块转换为 float8 进行训练的配置。
CastConfig	用于将单个张量可能转换为 float8 的配置
ScalingType
ScalingGranularity	定义转换为 float8 的缩放策略的粒度
precompute_float8_dynamic_scale_for_fsdp	为所有 float8 参数动态计算尺度。此函数应在优化器步骤之后运行。它执行一次 all-reduce 以计算所有 float8 权重的尺度。示例用法：model(input).sum().backward() optim.step() precompute_float8_dynamic_scale_for_fsdp(model)。