在 SageMaker 上追踪 LLM 微调 FLOPs，一次配置搞定 EU AI Act 合规审计

预计阅读时间：11 分钟

欧盟《人工智能法案》（EU AI Act）已经生效，其中一条容易被忽视的规则是：通用目的 AI（GPAI）模型如果训练或微调的计算量超过 10²⁵ FLOPs，将被归类为"系统性风险"模型，触发更严格的透明度与安全义务。 对在 AWS 上用 SageMaker 微调大模型的团队来说，问题很具体——你怎么知道自己的任务有没有踩线？又怎么向审计方证明？

开源工具 Fine-Tuning FLOPs Meter 正是为了解决这个问题：它嵌入 SageMaker 训练流程，自动统计 FLOPs，一个配置标志就能判定合规状态，并输出审计文档。下面拆解实操细节。

FLOPs 阈值为什么重要

EU AI Act 第 52 条和附件对 GPAI 模型设定了两档门槛：

类别	FLOPs 阈值	义务
标准 GPAI	< 10²⁵	基本透明度要求（训练数据摘要、版权合规声明）
系统性风险 GPAI	≥ 10²⁵	额外义务： adversarial 测试、风险评估、事件报告、严重事故通知

注意两点：第一，阈值针对的是单次训练或微调任务的累计 FLOPs，不是模型参数量；第二，微调（fine-tuning）同样计入——如果你在已有基座模型上做了大规模继续训练，FLOPs 累加后可能越过阈值。

这意味着每次微调任务结束后，你需要一个可信的 FLOPs 数字，而不是粗略估算。

Fine-Tuning FLOPs Meter 工作原理

Fine-Tuning FLOPs Meter 是一个轻量级 Python 工具包，核心思路是：

读取模型架构——从 HuggingFace config.json 提取层数、隐藏维度、注意力头数等参数。
结合训练超参——batch size、序列长度、训练步数、梯度累积步数。
按公式计算——使用 Chinchilla 论文中推导的 Transformer FLOPs 估算公式，区分前向传播与反向传播的乘数因子。
输出结果——生成 JSON 报告，包含总 FLOPs、是否触发系统性风险阈值、以及完整的参数快照供审计回溯。

它不需要修改训练脚本本身，而是作为 SageMaker 训练作业的一个附加组件运行。

在 SageMaker 上集成 FLOPs Meter

下面是一个完整的 SageMaker PyTorch 训练作业配置，集成了 FLOPs Meter。你可以直接改造后运行。

前置准备

# 安装 FLOPs Meter（本地测试用，SageMaker 作业中通过 requirements.txt 注入）
pip install fine-tuning-flops-meter

# 确认 AWS CLI 和 SageMaker Python SDK
pip install sagemaker
aws sts get-caller-identity  # 验证身份

训练脚本中嵌入 FLOPs 追踪

在你的 train.py 中加入以下逻辑。关键改动只有两处：训练前初始化 meter，训练后生成报告。

import json
import os
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset
from flops_meter import FLOPsMeter  # 核心导入

def main():
    model_id = os.environ.get("MODEL_ID", "meta-llama/Llama-3.1-8B")

    # ---- 正常加载模型和数据 ----
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype="auto",
        device_map="auto"
    )
    dataset = load_dataset("open-llm-leaderboard/mt-bench", split="train")

    # ---- 初始化 FLOPs Meter ----
    meter = FLOPsMeter(
        model_config=model.config,
        # 以下参数从环境变量读取，SageMaker 会自动注入
        batch_size=int(os.environ.get("TRAIN_BATCH_SIZE", 4)),
        seq_length=int(os.environ.get("MAX_SEQ_LENGTH", 2048)),
        gradient_accumulation_steps=int(os.environ.get("GRAD_ACCUM_STEPS", 2)),
        num_train_epochs=int(os.environ.get("NUM_EPOCHS", 1)),
        # 一个标志切换合规判定模式
        compliance_mode=os.environ.get("COMPLIANCE_MODE", "eu_ai_act")
    )

    training_args = TrainingArguments(
        output_dir="/opt/ml/model",
        per_device_train_batch_size=int(os.environ.get("TRAIN_BATCH_SIZE", 4)),
        gradient_accumulation_steps=int(os.environ.get("GRAD_ACCUM_STEPS", 2)),
        num_train_epochs=int(os.environ.get("NUM_EPOCHS", 1)),
        max_seq_length=int(os.environ.get("MAX_SEQ_LENGTH", 2048)),
        logging_steps=10,
        save_strategy="epoch",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
    )

    # ---- 训练前记录起始状态 ----
    meter.mark_start()

    trainer.train()

    # ---- 训练后计算 FLOPs 并生成报告 ----
    meter.mark_end()
    report = meter.generate_report()

    # 将报告写入 SageMaker 模型输出目录，随模型一起持久化
    report_path = "/opt/ml/model/flops_audit_report.json"
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2)

    print(f"Total FLOPs: {report['total_flops']:.2e}")
    print(f"Systemic risk threshold (10^25) exceeded: {report['exceeds_systemic_risk_threshold']}")
    print(f"Audit report saved to: {report_path}")

if __name__ == "__main__":
    main()

SageMaker 作业启动脚本

import sagemaker
from sagemaker.pytorch import PyTorch

sagemaker_session = sagemaker.session.Session()
role = sagemaker.get_execution_role()

# ---- 关键：通过 environment 传入合规配置 ----
env_config = {
    "MODEL_ID": "meta-llama/Llama-3.1-8B",
    "TRAIN_BATCH_SIZE": "4",
    "MAX_SEQ_LENGTH": "2048",
    "GRAD_ACCUM_STEPS": "2",
    "NUM_EPOCHS": "1",
    # 一个标志决定合规判定标准
    "COMPLIANCE_MODE": "eu_ai_act",
}

estimator = PyTorch(
    entry_point="train.py",
    source_dir="src/",           # 包含 train.py 和 requirements.txt
    instance_type="ml.p5.48xlarge",  # GPU 实例
    instance_count=1,
    role=role,
    framework_version="2.3",
    py_version="py311",
    environment=env_config,
    hyperparameters={
        # 超参也可以在这里定义，会覆盖 environment 中的同名值
    },
)

estimator.fit(wait=True)

# ---- 作业完成后下载审计报告 ----
model_artifacts = estimator.model_data  # S3 URI
# 用 AWS CLI 下载
# aws s3 cp {model_artifacts} ./model.tar.gz
# tar -xzf model.tar.gz flops_audit_report.json
# cat flops_audit_report.json

src/requirements.txt 中加入：

fine-tuning-flops-meter
transformers>=4.40
datasets

审计报告长什么样

运行完成后，flops_audit_report.json 大致如下：

{
  "total_flops": 3.72e+24,
  "exceeds_systemic_risk_threshold": false,
  "threshold_value": 1e25,
  "compliance_standard": "eu_ai_act_article_52",
  "model_config_snapshot": {
    "model_type": "llama",
    "num_hidden_layers": 32,
    "hidden_size": 4096,
    "num_attention_heads": 32,
    "intermediate_size": 14336,
    "num_key_value_heads": 8
  },
  "training_hyperparams": {
    "batch_size": 4,
    "seq_length": 2048,
    "gradient_accumulation_steps": 2,
    "num_train_epochs": 1,
    "effective_total_steps": 250
  },
  "calculation_method": "chinchilla_transformer_flops_estimate",
  "forward_pass_multiplier": 1.0,
  "backward_pass_multiplier": 2.0,
  "generated_at": "2025-01-15T10:23:00Z",
  "sagemaker_job_name": "llama31-8b-finetune-2025-01-15"
}

这份报告的价值在于：它把模型架构、训练超参、计算方法、合规判定全部打包在一起，审计方不需要追问"你这个数字怎么来的"。

合规标志的细节

COMPLIANCE_MODE 这个环境变量是整个流程中最省心的设计：

设为 eu_ai_act：自动对照 10²⁵ FLOPs 阈值判定，报告字段 exceeds_systemic_risk_threshold 为 true/false。
设为 none 或不设：只输出 FLOPs 数值，不做阈值判定——适合纯内部追踪场景。
未来如果其他监管框架出现不同阈值，只需切换此标志。

一个配置项替代了手动查表、手动比对、手动写结论的整个流程。

实操建议与风险边界

先算再跑。 在启动昂贵的 GPU 作业之前，用 FLOPs Meter 的 estimate_only 模式做一次预计算：

from flops_meter import FLOPsMeter
from transformers import AutoConfig

config = AutoConfig.from_pretrained("meta-llama/Llama-3.1-8B")
meter = FLOPsMeter(
    model_config=config,
    batch_size=4,
    seq_length=2048,
    gradient_accumulation_steps=2,
    num_train_epochs=3,  # 假设你打算跑 3 个 epoch
    compliance_mode="eu_ai_act"
)
estimate = meter.estimate_only()
print(f"预估 FLOPs: {estimate['total_flops']:.2e}")
print(f"是否会触发系统性风险: {estimate['exceeds_systemic_risk_threshold']}")

如果预估已经接近阈值，你有两个选择：缩短训练步数，或者提前准备系统性风险合规文档。

估算不是精确值。 FLOPs Meter 使用的是理论估算公式，不考虑 dropout、混合精度训练的实际加速、数据并行中的通信开销等因素。误差通常在 5-15% 范围内。对于合规判定来说，如果你的 FLOPs 在 9×10²⁴ 附近，估算的不确定性本身就意味着你应该按系统性风险来准备。

微调 ≠ 从头训练，但 FLOPs 照算。 即使你只微调了 1% 的参数，FLOPs 计算仍然基于完整模型的前向与反向传播（因为梯度仍然流过冻结层）。这是 EU AI Act 的立场，也是 FLOPs Meter 的计算方式。

合规检查清单

每次启动微调作业前，过一遍这张表：

[ ] 预估 FLOPs 是否低于 10²⁵？用 estimate_only 模式确认。
[ ] COMPLIANCE_MODE 是否设为 eu_ai_act？
[ ] 训练脚本是否调用了 meter.mark_start() 和 meter.mark_end()？
[ ] 审计报告是否写入 /opt/ml/model/ 目录，随模型产物持久化到 S3？
[ ] 如果 FLOPs 超阈值：是否准备了 adversarial 测试方案、风险评估文档、事件报告流程？
[ ] 报告中的 calculation_method 字段是否在内部合规文档中被引用为计算依据？

FLOPs 追踪不是 EU AI Act 合规的全部，但它是可量化、可自动化、可审计的第一步。把这一步嵌入 SageMaker 作业流水线，后续的透明度义务才有数据基础。