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)。