Function column_linear_forward

tests/unittest/_torch/multi_gpu/test_linear.py:112–147

(x, hidden_size, dtype, tensor_parallel_size,
                          tensor_parallel_rank, weights)

@torch.inference_mode
def column_linear_forward(x, hidden_size, dtype, tensor_parallel_size,
                          tensor_parallel_rank, weights):

    x = x.cuda()
    l0 = Linear(
        in_features=hidden_size,
        out_features=hidden_size,
        bias=False,
        dtype=dtype,
        mapping=Mapping(
            world_size=tensor_parallel_size,
            tp_size=tensor_parallel_size,
            rank=tensor_parallel_rank,
        ),
        tensor_parallel_mode=TensorParallelMode.COLUMN,
        gather_output=True,
    )
    l0.load_weights([dict(weight=weights[0])])
    l0.cuda()

    l0 = torch.compile(l0, fullgraph=True)
    output = l0.forward(x)

    # torch run
    l0 = nn.Linear(in_features=hidden_size,
                   out_features=hidden_size,
                   bias=False,
                   dtype=dtype)
    l0.weight.data.copy_(weights[0])
    l0.cuda()

    torch_output = l0.forward(x)

    # compare
    torch.cuda.synchronize()
    torch.testing.assert_close(output, torch_output)


@torch.inference_mode