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，表示不对优化器状态进行保存。

返回¶

无

代码示例¶

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

         