group_sharded_parallel¶

参数¶

• model (Layer) - 需要使用 group sharded 的模型。

• optimizer (Optimizer) - 需要使用 group sharded 的优化器。

• level (str) - 选择 group sharded 的级别，分别有'os','os_g','p_g_os'。

• scaler (GradScaler，可选) - 如果使用 AMP 混合精度，需要传入 GradScaler，默认为 None，表示不使用 GradScaler。

• group (Group，可选) - 工作的进程组编号，默认为 None，表示采用默认环境 Group。

• offload (bool，可选) - 是否使用 offload 缓存功能，默认为 False，表示不使用 offload 功能。

• sync_buffers (bool，可选) - 是否需要同步模型 buffers，一般在有注册模型 buffers 时才使用，默认为 False，表示不同步模型 buffers。

• buffer_max_size (int，可选) - 在'os_g'模式中会对梯度进行聚合，此选项指定聚合 buffer 的大小，指定越大则占用显存也越多，默认为 2**23，表示聚合 buffer 的维度为 2**23。

• segment_size (int，可选) - 在'p_g_os'模式中会对参数进行切分，此选项指定最小切分参数大小，默认为 2**20，表示最小被切分参数的维度为 2**20。

• sync_comm (bool，可选) - 在'p_g_os'模式中是否采用同步通信，默认为 False，表示使用异步通信流。

• dp_group (Group，可选) - 数据并行（data parallel）通信组，支持 sharding 的 stage2 和 stage3 和数据并行一起混合使用。

• exclude_layer (list，可选) - 在 sharding stage3 中，可以设置某些 layer 的参数不切分（通过 layer 的类型或 id 控制），例如：exclude_layer=["GroupNorm", id(model.gpt.linear)]。

返回¶

group sharded 配置后的 model，optimizer 和 scaler

代码示例¶

>>> import paddle
>>> from paddle.nn import Linear
>>> from paddle.distributed import fleet
>>> from paddle.distributed.sharding import group_sharded_parallel

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