一、optimizer.py深度分析1. 架构设计┌─────────────────────────────────────────────────────────────────┐ │ 优化器配置系统 │ ├─────────────────────────────────────────────────────────────────┤ │ LRScheduleConfig (Protocol) │ │ ├── CosineDecaySchedule │ │ └── RsqrtDecaySchedule │ ├─────────────────────────────────────────────────────────────────┤ │ OptimizerConfig (Protocol) │ │ ├── AdamW │ │ └── SGD │ ├─────────────────────────────────────────────────────────────────┤ │ create_optimizer() → optax.GradientTransformation │ └─────────────────────────────────────────────────────────────────┘2. 学习率调度器CosineDecaySchedule余弦衰减# optimizer.py:16-31dataclasses.dataclass(frozenTrue)classCosineDecaySchedule(LRScheduleConfig):warmup_steps:int1_000peak_lr:float2.5e-5decay_steps:int30_000decay_lr:float2.5e-6defcreate(self)-optax.Schedule:returnoptax.warmup_cosine_decay_schedule(init_valueself.peak_lr/(self.warmup_steps1),peak_valueself.peak_lr,warmup_stepsself.warmup_steps,decay_stepsself.decay_steps,end_valueself.decay_lr,)学习率曲线学习率 ^ │ peak_lr │ /\ │ / \ │ / \ │ / \ │ / \ │ / \ └──┴────────────┴──→ 训练步数 warmup decay init_value decay_lr参数说明参数默认值作用warmup_steps1000线性预热步数peak_lr2.5e-5峰值学习率decay_steps30000衰减总步数decay_lr2.5e-6最终学习率峰值的 10%设计意图Warmup从很小的学习率开始逐渐增加到峰值稳定训练初期的梯度Cosine Decay后期学习率平滑衰减精细调整参数避免震荡RsqrtDecaySchedule平方根倒数衰减# optimizer.py:34-53dataclasses.dataclass(frozenTrue)classRsqrtDecaySchedule(LRScheduleConfig):warmup_steps:int1_000peak_lr:float5e-5timescale:float10_000defcreate(self)-optax.Schedule:returnoptax.join_schedules([optax.linear_schedule(init_valueself.peak_lr/(self.warmup_steps1),end_valueself.peak_lr,transition_stepsself.warmup_steps,),lambdastep:self.peak_lr/jnp.sqrt((self.timescalestep)/self.timescale),],[self.warmup_steps],)学习率曲线学习率 ^ │ peak_lr │ /\ │ / \ │ / \ │ / \ │ / \ │ / \ └──┴────────────┴──→ 训练步数 warmup 1/√t 衰减衰减公式lr(t) peak_lr / √((timescale t) / timescale)设计意图衰减速度比余弦衰减慢适合需要更长训练时间的场景理论基础梯度噪声的方差按 1/√t 衰减3. 优化器AdamW# optimizer.py:65-85dataclasses.dataclass(frozenTrue)classAdamW(OptimizerConfig):b1:float0.9b2:float0.95eps:float1e-8weight_decay:float1e-10clip_gradient_norm:float1.0defcreate(self,lr,weight_decay_maskNone):txoptax.adamw(lr,b1self.b1,b2self.b2,epsself.eps,weight_decayself.weight_decay,maskweight_decay_mask)returnoptax.chain(optax.clip_by_global_norm(self.clip_gradient_norm),tx)关键设计参数值原因b10.9一阶动量衰减率标准设置b20.95二阶动量衰减率比标准 0.999 小更快适应梯度变化weight_decay1e-10极小值注释说设为 0 会导致 OOM 错误clip_gradient_norm1.0梯度裁剪防止梯度爆炸优化器链optax.chain( optax.clip_by_global_norm(1.0), # 先裁剪梯度 optax.adamw(...) # 再应用 AdamW 更新 )设计意图先裁剪梯度防止爆炸再应用 AdamW 更新权重衰减通过weight_decay_mask控制项目中未使用SGD# optimizer.py:88-102dataclasses.dataclass(frozenTrue)classSGD(OptimizerConfig):lr:float5e-5momentum:float0.9nesterov:boolFalsedefcreate(self,lr,weight_decay_maskNone):assertweight_decay_maskisNone,Weight decay is not supported for SGDreturnoptax.sgd(lr,momentumself.momentum,nesterovself.nesterov)设计意图简单的 SGD 实现支持动量加速不支持权重衰减需要单独实现4. 创建函数# optimizer.py:105-109defcreate_optimizer(optimizer:OptimizerConfig,lr_schedule:LRScheduleConfig,weight_decay_mask:at.PyTree|NoneNone)-optax.GradientTransformation:lrlr_schedule.create()returnoptimizer.create(lr,weight_decay_maskweight_decay_mask)作用将配置对象转换为 optax 的GradientTransformation供训练循环使用。二、sharding.py深度分析1. 架构设计┌─────────────────────────────────────────────────────────────────┐ │ 分布式训练系统 │ ├─────────────────────────────────────────────────────────────────┤ │ 常量定义 │ │ ├── BATCH_AXIS batch │ │ ├── FSDP_AXIS fsdp │ │ └── DATA_AXIS (batch, fsdp) │ ├─────────────────────────────────────────────────────────────────┤ │ _MeshState (全局状态) │ │ └── active_mesh: jax.sharding.Mesh | None │ ├─────────────────────────────────────────────────────────────────┤ │ make_mesh() → jax.sharding.Mesh │ │ set_mesh() (context manager) │ │ activation_sharding_constraint() │ │ fsdp_sharding() │ └─────────────────────────────────────────────────────────────────┘2. 分布式训练概念FSDPFully Sharded Data Parallel传统数据并行: 每个 GPU 拥有完整的模型副本 梯度在 GPU 间同步 显存占用: O(N × model_size) FSDP: 模型权重在 GPU 间分片 每个 GPU 只拥有部分权重 显存占用: O(model_size)优势可以训练更大的模型。3. 核心组件轴定义# sharding.py:7-10BATCH_AXISbatchFSDP_AXISfsdpDATA_AXIS(BATCH_AXIS,FSDP_AXIS)说明BATCH_AXIS数据批次维度数据在多个 GPU 上并行计算FSDP_AXIS模型权重分片维度权重在多个 GPU 上分布存储DATA_AXIS数据同时在 batch 和 fsdp 轴上分布_MeshState全局状态# sharding.py:13-14class_MeshState:active_mesh:jax.sharding.Mesh|NoneNone设计意图全局存储当前激活的设备网格避免在整个模块树中传递 mesh 参数通过 context manager 管理生命周期make_mesh创建网格# sharding.py:17-23defmake_mesh(num_fsdp_devices:int)-jax.sharding.Mesh:ifjax.device_count()%num_fsdp_devices!0:raiseValueError(...)mesh_shape(jax.device_count()//num_fsdp_devices,num_fsdp_devices)returnjax.make_mesh(mesh_shape,(BATCH_AXIS,FSDP_AXIS))计算逻辑输入示例jax.device_count()88 个 GPUnum_fsdp_devices4每 4 个 GPU 分片一个模型输出mesh_shape (8 // 4, 4) (2, 4) mesh jax.make_mesh((2, 4), (batch, fsdp))含义2 个 batch 组每组处理不同的数据批次每个 batch 组内有 4 个 GPU 分片模型权重set_mesh上下文管理器# sharding.py:26-37contextlib.contextmanagerdefset_mesh(mesh:jax.sharding.Mesh):if_MeshState.active_meshisnotNone:raiseValueError(Cannot nest set_mesh context managers.)_MeshState.active_meshmeshtry:yieldfinally:_MeshState.active_meshNone设计意图确保同一时间只有一个活跃的 mesh自动清理状态避免内存泄漏简化 mesh 的使用方式activation_sharding_constraint激活函数分片约束# sharding.py:40-45defactivation_sharding_constraint(pytree):if_MeshState.active_meshisNone:returnpytreereturnjax.lax.with_sharding_constraint(pytree,jax.sharding.NamedSharding(_MeshState.active_mesh,jax.sharding.PartitionSpec(DATA_AXIS)))作用对激活函数输出施加分片约束数据同时在 batch 和 fsdp 轴上分布优化数据通信效率使用场景gemma.py:294xssharding.activation_sharding_constraint(xs)fsdp_shardingFSDP 权重分片# sharding.py:48-100deffsdp_sharding(pytree,mesh,*,min_size_mbytes4,logFalse):min_size_bytesmin_size_mbytes*2**20def_shard_arr(kp,array):ifmesh.shape[FSDP_AXIS]1:returnjax.sharding.NamedSharding(mesh,jax.sharding.PartitionSpec())iflen(array.shape)2:returnjax.sharding.NamedSharding(mesh,jax.sharding.PartitionSpec())ifarr_sizemin_size_bytes:returnjax.sharding.NamedSharding(mesh,jax.sharding.PartitionSpec())axesnp.argsort(array.shape)[::-1]spec[None]*len(axes)foriinaxes:ifarray.shape[i]%mesh.shape[FSDP_AXIS]0:spec[i]FSDP_AXISreturnjax.sharding.NamedSharding(mesh,jax.sharding.PartitionSpec(*spec))returnjax.sharding.NamedSharding(mesh,jax.sharding.PartitionSpec())returnjax.tree_util.tree_map_with_path(_shard_arr,pytree)分片策略条件行为mesh.shape[FSDP_AXIS] 1不使用 FSDP复制所有参数数组维度 2复制标量和向量数组大小 4MB复制小数组分片开销大存在可分轴沿最大的可分轴分片无合适轴复制设计意图只对大型矩阵分片节省通信开销沿最大的可分轴分片最大化并行度小参数复制避免分片开销4. 使用流程# 1. 创建设备网格meshmake_mesh(num_fsdp_devices4)# 2. 设置全局网格withset_mesh(mesh):# 3. 对模型参数施加分片约束sharded_paramsfsdp_sharding(params,mesh)# 4. 在模型前向/反向传播中自动应用分片lossmodel.compute_loss(sharded_params,...)三、两个文件的关系┌─────────────────────────────────────────────────────────────────┐ │ 训练系统架构 │ ├─────────────────────────────────────────────────────────────────┤ │ 训练循环 │ │ ┌──────────────────┐ ┌──────────────────┐ │ │ │ optimizer.py │ │ sharding.py │ │ │ │ │ │ │ │ │ │ LRScheduleConfig │ │ jax.sharding.Mesh│ │ │ │ OptimizerConfig │ │ FSDP 分片 │ │ │ │ create_optimizer │ │ 激活函数约束 │ │ │ │ │ │ │ │ │ └────────┬─────────┘ └────────┬─────────┘ │ │ │ │ │ │ ▼ ▼ │ │ ┌──────────────────────────────────────────────┐ │ │ │ 模型训练循环 │ │ │ │ 前向传播 → 损失计算 → 反向传播 → 参数更新 │ │ │ │ (分片执行) (分片执行) (分片执行) │ │ │ └──────────────────────────────────────────────┘ │ └─────────────────────────────────────────────────────────────────┘协同工作optimizer.py负责参数更新策略sharding.py负责分布式训练的硬件资源管理两者结合实现高效的大规模模型训练四、总结optimizer.py组件作用设计亮点LRScheduleConfig学习率调度协议支持多种调度策略CosineDecaySchedule余弦衰减平滑衰减适合大多数场景RsqrtDecaySchedule平方根衰减理论保证适合长训练AdamW优化器梯度裁剪 权重衰减防止 OOMSGD优化器简单实现备用选项sharding.py组件作用设计亮点make_mesh创建设备网格自动计算网格形状set_mesh上下文管理器全局状态管理避免嵌套activation_sharding_constraint激活函数分片自动应用分片约束fsdp_sharding权重分片智能分片策略避免小参数分片核心价值optimizer.py实现了灵活的优化器配置支持多种学习率调度和优化算法sharding.py实现了高效的分布式训练支持 FSDP 权重分片降低显存占用