save_group_sharded_model

paddle.distributed.sharding. save_group_sharded_model ( model, output, optimizer=None ) [源代码]

使用 save_group_sharded_model 可以对 group_sharded_parallel 配置后的模型和优化器状态进行保存。

注解

此处需要注意,使用 save_group_sharded_model 保存模型,再次 load 时需要在调用 group_sharded_parallel 前对 model 和 optimizer 进行 set_state_dict。

参数

  • model (Layer) - 使用 group_sharded_parallel 配置后的模型。

  • output (str) - 输出保存模型和优化器的文件夹路径。

  • optimizer (Optimizer,可选) - 使用 group_sharded_parallel 配置后的优化器,默认为 None,表示不对优化器状态进行保存。

返回

代码示例

# required: distributed
import paddle
from paddle.nn import Linear
from paddle.distributed import fleet
from paddle.distributed.sharding import group_sharded_parallel, save_group_sharded_model

fleet.init(is_collective=True)
group = paddle.distributed.new_group([0, 1])
model = Linear(1000, 1000)

clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=1.0)
optimizer = paddle.optimizer.AdamW(learning_rate=0.001, parameters=model.parameters(), weight_decay=0.00001, grad_clip=clip)

# wrap sharding model, optimizer and scaler
model, optimizer, scaler = group_sharded_parallel(model, optimizer, "p_g", scaler=scaler)

img, label = data
label.stop_gradient = True
img.stop_gradient = True

out = model(img)
loss = paddle.nn.functional.cross_entropy(input=out, label=label)

loss.backward()
optimizer.step()
optimizer.clear_grad()

# save model and optimizer state_dict
save_group_sharded_model(model, optimizer, output=output_dir)