(model_parallel_size)
| 185 | |
| 186 | |
| 187 | def test_column_parallel_linear(model_parallel_size): |
| 188 | |
| 189 | mpu.initialize_model_parallel(model_parallel_size) |
| 190 | if torch.distributed.get_rank() == 0: |
| 191 | print('> testing ColumnParallelLinear with model parallel ' |
| 192 | 'size: {}'.format(model_parallel_size)) |
| 193 | model_parallel_size = mpu.get_model_parallel_world_size() |
| 194 | |
| 195 | seed = 12345 |
| 196 | set_random_seed(seed) |
| 197 | input_size_coeff = 13 |
| 198 | input_size = input_size_coeff * model_parallel_size |
| 199 | output_size_coeff = 17 |
| 200 | output_size = output_size_coeff * model_parallel_size |
| 201 | batch_size = 7 |
| 202 | |
| 203 | # Network |
| 204 | identity_layer = IdentityLayer2D(batch_size, input_size).cuda() |
| 205 | linear_layer = mpu.ColumnParallelLinear( |
| 206 | input_size, output_size, keep_master_weight_for_test=True).cuda() |
| 207 | loss_weight = torch.randn([batch_size, output_size]).cuda() |
| 208 | # Forward |
| 209 | input_ = identity_layer() |
| 210 | output = linear_layer(input_) |
| 211 | loss = torch.mul(output, loss_weight).sum() |
| 212 | # Backward |
| 213 | loss.backward() |
| 214 | |
| 215 | # Values. |
| 216 | dLdY = loss_weight |
| 217 | X = identity_layer.weight |
| 218 | A = linear_layer.master_weight.cuda() |
| 219 | dLdA = torch.matmul(dLdY.t(), X) |
| 220 | dLdb = torch.matmul(torch.ones(batch_size, 1).cuda().t(), dLdY).view(-1) |
| 221 | dLdX = torch.matmul(dLdY, A) |
| 222 | |
| 223 | rank = mpu.get_model_parallel_rank() |
| 224 | my_dLdA = torch.split(dLdA, output_size_coeff, |
| 225 | dim=0)[rank].contiguous().clone() |
| 226 | error = my_dLdA.sub(linear_layer.weight.grad).abs().max() |
| 227 | torch.distributed.barrier() |
| 228 | print(' error in dLdA on global rank {}: {}'.format( |
| 229 | torch.distributed.get_rank(), error)) |
| 230 | assert error < 1.0e-6 |
| 231 | |
| 232 | my_dLdb = torch.split(dLdb, output_size_coeff, |
| 233 | dim=0)[rank].contiguous().clone() |
| 234 | error = my_dLdb.sub(linear_layer.bias.grad).abs().max() |
| 235 | torch.distributed.barrier() |
| 236 | print(' error in dLdb on global rank {}: {}'.format( |
| 237 | torch.distributed.get_rank(), error)) |
| 238 | assert error < 1.0e-6 |
| 239 | |
| 240 | error = dLdX.sub(identity_layer.weight.grad).abs().max() |
| 241 | torch.distributed.barrier() |
| 242 | print(' error in dLdX on global rank {}: {}'.format( |
| 243 | torch.distributed.get_rank(), error)) |
| 244 | assert error < 1.0e-6 |
no test coverage detected