3D Tensor Parallelism

[conceptofmind]

Hi,

When defining layers for 3D Tensor Parallelism, Colossal-AI will automatically apply 3D parallelism to all the layers from colossalai.nn.

For example:

CONFIG = dict(parallel=dict(
    data=1,
    pipeline=1,
    tensor=dict(size=8, mode='3d'),
))

dense_1 = colossalai.nn.Linear(in_features, out_features)

In the documentation there are also specific layers defined for 3D Tensor Parallelism:

dense_1 = colossalai.nn.layer.parallel_3d.Linear3D(in_features, out_features)

Are these two examples equivalent for 3D Tensor Parallelism? Can they be used interchangeably?

I greatly appreciate the help.

Thank you,

Enrico

[kurisusnowdeng]

Sure. It is equivalent to using colossalai.nn.Linear when 3D parallelism is configured.

[conceptofmind]

Hi @kurisusnowdeng ,

I appreciate the response.

I just wanted to confirm that they could be used in a similar fashion.

Thank you for the help,

Enrico