[AutoTP] Make AutoTP work when num_heads not divisible by number of workers

[delock]

Currently AutoTP will assert when num_heads are not divisible by number of workers. However in some situation this might be what user intend to do. i.e. having three compute device for model with 16 heads and want to put all three compute device to work. In this situation, each worker will process 5 or 6 heads, which is still better than run the workload on two compute device and each worker process 8 heads and leave the third device idle.

This PR distribute attention heads to each worker as even as possible, and shard hidden_size according to this distribution, this allows AutoTP run OOB even when number heads are not divisible by number of devices installed on the system.

[mrwyattii]

Can we add a unit test to verify for an odd number of devices? Perhaps extend this test class:

DeepSpeed/tests/unit/inference/test_inference.py

Line 489 in 777ae39

class TestAutoTensorParallelism(DistributedTest):

to add something like:

@pytest.mark.world_size(3)
def test_odd_world_size(
        self,
        model_w_task,
        query,
        inf_kwargs,
        assert_fn,
        dtype,
    ):
        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
        if invalid_test_msg:
            pytest.skip(invalid_test_msg)

        model, task = model_w_task
        local_rank = int(os.getenv("LOCAL_RANK", "0"))
        world_size = int(os.getenv("WORLD_SIZE", "2"))

        pipe = pipeline(task, model=model, device=torch.device("cpu"), framework="pt")
        bs_output = pipe(query, **inf_kwargs)

        pipe.model = deepspeed.init_inference(pipe.model, mp_size=world_size, dtype=dtype)
        # Switch device to GPU so that input tensors are not on CPU
        pipe.device = torch.device(get_accelerator().device_name(local_rank))
        ds_output = pipe(query, **inf_kwargs)

        print(local_rank, "baseline", bs_output)
        print(local_rank, "deepspeed", ds_output)
        assert assert_fn(bs_output, ds_output)

[delock]

@mrwyattii Test added. There is a result mismatch assertion in the test and I can also reproduce this assertion with CPU+BF16. Will need sometime to debug this issue.

Can we add a unit test to verify for an odd number of devices? Perhaps extend this test class:

DeepSpeed/tests/unit/inference/test_inference.py

Line 489 in 777ae39

class TestAutoTensorParallelism(DistributedTest):

to add something like:

@pytest.mark.world_size(3)
def test_odd_world_size(
        self,
        model_w_task,
        query,
        inf_kwargs,
        assert_fn,
        dtype,
    ):
        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
        if invalid_test_msg:
            pytest.skip(invalid_test_msg)

        model, task = model_w_task
        local_rank = int(os.getenv("LOCAL_RANK", "0"))
        world_size = int(os.getenv("WORLD_SIZE", "2"))

        pipe = pipeline(task, model=model, device=torch.device("cpu"), framework="pt")
        bs_output = pipe(query, **inf_kwargs)

        pipe.model = deepspeed.init_inference(pipe.model, mp_size=world_size, dtype=dtype)
        # Switch device to GPU so that input tensors are not on CPU
        pipe.device = torch.device(get_accelerator().device_name(local_rank))
        ds_output = pipe(query, **inf_kwargs)

        print(local_rank, "baseline", bs_output)
        print(local_rank, "deepspeed", ds_output)
        assert assert_fn(bs_output, ds_output)

[delock]

@mrwyattii @molly-smith I have identified the issue for result mismatch and fixed. Can you help restart workflow? Thanks!

@mrwyattii Test added. There is a result mismatch assertion in the test and I can also reproduce this assertion with CPU+BF16. Will need sometime to debug this issue.

Can we add a unit test to verify for an odd number of devices? Perhaps extend this test class:

DeepSpeed/tests/unit/inference/test_inference.py

Line 489 in 777ae39

class TestAutoTensorParallelism(DistributedTest):

to add something like:

@pytest.mark.world_size(3)
def test_odd_world_size(
        self,
        model_w_task,
        query,
        inf_kwargs,
        assert_fn,
        dtype,
    ):
        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
        if invalid_test_msg:
            pytest.skip(invalid_test_msg)

        model, task = model_w_task
        local_rank = int(os.getenv("LOCAL_RANK", "0"))
        world_size = int(os.getenv("WORLD_SIZE", "2"))

        pipe = pipeline(task, model=model, device=torch.device("cpu"), framework="pt")
        bs_output = pipe(query, **inf_kwargs)

        pipe.model = deepspeed.init_inference(pipe.model, mp_size=world_size, dtype=dtype)
        # Switch device to GPU so that input tensors are not on CPU
        pipe.device = torch.device(get_accelerator().device_name(local_rank))
        ds_output = pipe(query, **inf_kwargs)

        print(local_rank, "baseline", bs_output)
        print(local_rank, "deepspeed", ds_output)
        assert assert_fn(bs_output, ds_output)

[delock]

Hi @mrwyattii @molly-smith The test failure is fixed. Can you help restart CI workflow? Thanks!

[mrwyattii]

@delock approved the PR, but there is a merge conflict. Can you resolve that? The PR will auto-merge after!

[tjruwase]

@delock, can you please help with the merge conflict?

[delock]

@mrwyattii @tjruwase the conflict is resolved, thanks!

[delock]

Conflict with lm_head parallelism resolved, and add uneven sharding support for lm_head parallel.

[delock]

Hi @mrwyattii @tjruwase , the recent merge conflict had been resolved, and we also support uneven sharding of lm_head parallel. Can you take a quick look whether it can be put into merge queue? Thanks!