[Bug] Get crashed when change batchsize in bmk_comm_latency_multiserver.py

[Zhangmj0621]

When we change bsz in bmk_comm_latency_multiserver.py, we get error logs like below.
RuntimeError: [15:17:41] /infrawaves/StepMesh/fserver/csrc/./public.hpp:103: Check failed: (tensors.size()) == (reqmeta.pull_tensors.size())
After deep dive into the code, we found that the hard code below only use the first tensor in input_tensors.

elif is_server:
    ret_buffer = torch.rand([65535, dim], dtype=torch.bfloat16, device='cuda')
    count = 0
    f.barrier(True, False)
    def server():
        global count
        iter_count = 0
        while True:
            batches = f.get_batch()
            if len(batches) != 0:
                iter_count += 1
                # hard code
                - recv_tensor_list = [batches[i][1][0] for i in range(worker_count)]
                + recv_tensor_list = [batches[i][1] for i in range(worker_count)]
                comm_id_list = [batches[i][0] for i in range(worker_count)]

                f.respond_vec(ret_buffer, recv_tensor_list, comm_id_list)
                if iter_count == num_iters:
                    break
    server()

f.stop()

Besides change the hard code above, we also change the related code in fserver/csrc/public.hpp as follows:

void respond_vec(torch::Tensor& ret_buffer,
                 std::vector<std::vector<torch::Tensor>>& tensors_vec,
                 std::vector<uint64_t>& handler_vec) {
  PS_CHECK_EQ(tensors_vec.size(), handler_vec.size());
  for (size_t i = 0; i < handler_vec.size(); i++) {
    std::vector<torch::Tensor> sliced_buffer_list;
    int64_t tensor_shape_0 = tensors_vec[i][0].size(0);
    for (int j = 0; j < tensors_vec[i].size() - 1; j++) {
      sliced_buffer_list.push_back(
          ret_buffer.slice(0, j * tensor_shape_0, tensor_shape_0));
    }
    //std::vector<torch::Tensor> sliced_buffer_list = {
    //    ret_buffer.slice(0, 0, tensor_shape_0)
    //};
    respond(sliced_buffer_list, handler_vec[i], i == 0);
  }
}

Unfortunatly, we see further error as follows.

[01:55:14] server /gpfs/Stepmesh/include/dmlc/logging.h:301: [01:55:14] /infrawaves/StepMesh/src/./rdma_van.h:875: Check failed: (temp_mr) != (mem_mr_.end()) 
Stack trace returned 6 entries:
[bt] (0) /gpfs/Stepmesh/fserver_lib.cpython-310-x86_64-linux-gnu.so(+0x5422f) [0x7f897146822f]
[bt] (1) /gpfs/Stepmesh/fserver_lib.cpython-310-x86_64-linux-gnu.so(+0x54563) [0x7f8971468563]
[bt] (2) /gpfs/Stepmesh/fserver_lib.cpython-310-x86_64-linux-gnu.so(ps::RDMAVan::PollCQ()+0x1a81) [0x7f89714d8951]
[bt] (3) /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xdc253) [0x7f8a8b6b0253]
[bt] (4) /usr/lib/x86_64-linux-gnu/libc.so.6(+0x94ac3) [0x7f8aeb334ac3]
[bt] (5) /usr/lib/x86_64-linux-gnu/libc.so.6(+0x126850) [0x7f8aeb3c6850]

So I wonder if you have encounter the question below and if any solutions to this. I will highly appreciate it if any feedbacks from you.

[niehao100]

Hi,
This is a constraint for the respond tensor. It request the size of respond tensor is the same as the tensor set for push_pull api( pull api), due to the mr limitation of rdma.

So please print the shape of them at both server/worker side and make sure they are the same.

[niehao100]

The type of batches is std::vector<std::vector<std::tuple<uint64_t, std::vectortorch::Tensor, std::vector<uint64_t>>>, so I think batches[i][0] is a tuple, this maybe the root cause. Hope this can help u ☺

[Zhangmj0621]

batches[i][0] is original code. The type of batches is std::vector<ServerDataBatch> which is indeed std::vector<std::tuple<uint64_t, std::vector<torch::Tensor>, std::vector<uint64_t>>>

[niehao100]

Yes you are right

batches[i][0] is original code. The type of batches is std::vector which is indeed std::vector<std::tuple<uint64_t, std::vectortorch::Tensor, std::vector<uint64_t>>>

[niehao100]

The batches is the vector of push request of all workers.
And the tensors vector in ServerDataBatch is the pushed tensors from worker.

In bmk_comm_latency_multiserver.py, the pushed tensor list inp_tensors_buffers[idx % 3] contains only one data tensor and another expert_token_cnt_buffer as scale tensor. The scale tensor doesn't match a pull tensor, please make sure it is skipped in respond vec.

In addition ,the respond api can be called directly in python, and you can implement the respond_vec function in python too.

[Zhangmj0621]

The batches is the vector of push request of all workers. And the tensors vector in ServerDataBatch is the pushed tensors from worker.

In bmk_comm_latency_multiserver.py, the pushed tensor list inp_tensors_buffers[idx % 3] contains only one data tensor and another expert_token_cnt_buffer as scale tensor. The scale tensor doesn't match a pull tensor, please make sure it is skipped in respond vec.

In addition ,the respond api can be called directly in python, and you can implement the respond_vec function in python too.

So you mean StepMesh only supports batchsize=1 now and Implementing support for flexible batch sizes would indeed require additional development effort.

[niehao100]

The batches is the vector of push request of all workers. And the tensors vector in ServerDataBatch is the pushed tensors from worker.
In bmk_comm_latency_multiserver.py, the pushed tensor list inp_tensors_buffers[idx % 3] contains only one data tensor and another expert_token_cnt_buffer as scale tensor. The scale tensor doesn't match a pull tensor, please make sure it is skipped in respond vec.
In addition ,the respond api can be called directly in python, and you can implement the respond_vec function in python too.

So you mean StepMesh only supports batchsize=1 now and Implementing support for flexible batch sizes would indeed require additional development effort.

StepMesh support larger batch size. The crash is only caused by the design in bmk_comm_latency_multiserver.py. Change the size of tokens_buffers and o_tensors by bsz, then you can reponse more tensors with the respond api at server side.

[Zhangmj0621]

However, in bmk_comm_latency_multiserver.py, it seems the size of tokens_buffers and o_tensors will scale by bsz, which means inp_tensors_buffers[idx % 3] can have serveral data tensors already. So I can't understand what is the key limitations that restrict large batch size.

for mb in range(3):
    expert_token_cnt_buffer = torch.tensor([num_token for _ in range(bsz)], dtype=torch.int32, device=f'cuda:{gpu}')
    tokens_buffers = [
        torch.rand([num_token, dim], dtype=torch.bfloat16, device=f'cuda:{gpu}') for _ in range(bsz)
    ]
    tokens_buffers.append(expert_token_cnt_buffer)
    inp_tensors_buffers.append(tokens_buffers)
    inp_tensors_keys.append([gen_push_key(i, mb) for i in range(len(tokens_buffers))])

    o_tensors = []
    for _ in range(server_count ):
        o_tensors += [torch.rand([num_token, dim], dtype=torch.bfloat16, device=f'cuda:{gpu}') for _ in range(bsz)]
    out_tensors_buffers.append(o_tensors)

    out_tensors_keys.append([gen_pull_key(i, mb) for i in range(len(o_tensors))])

[niehao100]

Hi @Zhangmj0621 . This new version of test_fserver.py push 3 tensors and pull 2 tenors, hope this example will help.

import torch, os
import time
import fserver_lib as f
is_worker = os.environ.get('DMLC_ROLE') == 'worker'
is_server = os.environ.get('DMLC_ROLE') == 'server'

f.init()

if is_worker:
    gpu = os.environ.get('STEPMESH_GPU')
    push_tensors = [
        torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}'),
        torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}'),
        torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}'),
    ]
    pull_tensors = [
        torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}'),
        torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}')
    ]
    handler = f.push_pull(
        push_tensors, 
        [i for i in range(len(push_tensors))], 
        pull_tensors, 
        [i for i in range(len(pull_tensors))],
    )
    f.wait(handler)
    print(torch.allclose(push_tensors[0], pull_tensors[0]))
    print(torch.allclose(push_tensors[1], pull_tensors[1]))
    print("worker test done")

elif is_server:
    gpu = os.environ.get('STEPMESH_GPU')
    torch.set_default_device('cuda:{}'.format(gpu))
    res = []
    while len(res) == 0:
        time.sleep(1)
        res = f.get_batch()
    print(res)
    for r in res:
        comm_id, batch, _ = r
        f.respond([batch[0],batch[1]], comm_id)

f.stop()

[niehao100]

However, in bmk_comm_latency_multiserver.py, it seems the size of tokens_buffers and o_tensors will scale by bsz, which means inp_tensors_buffers[idx % 3] can have serveral data tensors already. So I can't understand what is the key limitations that restrict large batch size.

for mb in range(3):
    expert_token_cnt_buffer = torch.tensor([num_token for _ in range(bsz)], dtype=torch.int32, device=f'cuda:{gpu}')
    tokens_buffers = [
        torch.rand([num_token, dim], dtype=torch.bfloat16, device=f'cuda:{gpu}') for _ in range(bsz)
    ]
    tokens_buffers.append(expert_token_cnt_buffer)
    inp_tensors_buffers.append(tokens_buffers)
    inp_tensors_keys.append([gen_push_key(i, mb) for i in range(len(tokens_buffers))])

    o_tensors = []
    for _ in range(server_count ):
        o_tensors += [torch.rand([num_token, dim], dtype=torch.bfloat16, device=f'cuda:{gpu}') for _ in range(bsz)]
    out_tensors_buffers.append(o_tensors)

    out_tensors_keys.append([gen_pull_key(i, mb) for i in range(len(o_tensors))])

We use mb and idx%s to simulate the three stage pipeline in tech report. Batch them together is ok for communication.

[Zhangmj0621]

Hi @Zhangmj0621 . This new version of test_fserver.py push 3 tensors and pull 2 tenors, hope this example will help.

import torch, os
import time
import fserver_lib as f
is_worker = os.environ.get('DMLC_ROLE') == 'worker'
is_server = os.environ.get('DMLC_ROLE') == 'server'

f.init()

if is_worker:
gpu = os.environ.get('STEPMESH_GPU')
push_tensors = [
torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}'),
torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}'),
torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}'),
]
pull_tensors = [
torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}'),
torch.rand([1, 8192], dtype=torch.float32, device=f'cuda:{gpu}')
]
handler = f.push_pull(
push_tensors,
[i for i in range(len(push_tensors))],
pull_tensors,
[i for i in range(len(pull_tensors))],
)
f.wait(handler)
print(torch.allclose(push_tensors[0], pull_tensors[0]))
print(torch.allclose(push_tensors[1], pull_tensors[1]))
print("worker test done")

elif is_server:
gpu = os.environ.get('STEPMESH_GPU')
torch.set_default_device('cuda:{}'.format(gpu))
res = []
while len(res) == 0:
time.sleep(1)
res = f.get_batch()
print(res)
for r in res:
comm_id, batch, _ = r
f.respond([batch[0],batch[1]], comm_id)

f.stop()

When I test this benchmark script, I found that the e2e latency is 120ms. However, when I change the push tensor num from 3 to 1, I observed that the latency scales to 600ms. Do you know any potential reasons causing it and why the time counted here is much larger than the time provided in benchmark bmk_comm_latency_multiserver.py?

[niehao100]

Hi @Zhangmj0621 The reason is this script runs push_pull only once. The tenors data requires ibv_reg_mr when it is used by RNIC at the first time, and the ibv_reg_mr time consumption is unstable.

[Zhangmj0621]

Hi @Zhangmj0621 The reason is this script runs push_pull only once. The tenors data requires ibv_reg_mr when it is used by RNIC at the first time, and the ibv_reg_mr time consumption is unstable.

I get the point why this script is slow. But why is push 3 tensor quicker than push 1 tensor? They both run push_pull only once.

[niehao100]

Hi @Zhangmj0621 I try to reproduce on our testbed, the following push one and two tensor and pull only one. The time is calculated with perf_counter. The time consumption is close, and I think the network consumption is not the main part.

    start = time.perf_counter()
    handler = f.push_pull(
        push_tensors, 
        [i for i in range(len(push_tensors))], 
        pull_tensors, 
        [i for i in range(len(pull_tensors))],
    )
    f.wait(handler)
    # perf_counter as ms
    print(f"time: {(time.perf_counter() - start)*1e3:.2f} ms")

-> # ROLE=joint RNIC=brainpf0  BIN=test_multi_push bash tests/fserver/run_single_gpu.sh
kill all testing process of ps lite for user root
SCHEDULER_IP: 10.195.170.30
time: 953.23 ms
True
push_tensors: [tensor([[0.8306, 0.2724, 0.1026,  ..., 0.4611, 0.7460, 0.4321]],
       device='cuda:0'), tensor([[0.3074, 0.9407, 0.7333,  ..., 0.3799, 0.8943, 0.5130]],
       device='cuda:0')]
pull_tensors: [tensor([[0.8306, 0.2724, 0.1026,  ..., 0.4611, 0.7460, 0.4321]],
       device='cuda:0')]
worker test done
kill all testing process of ps lite for user root

-> # ROLE=joint RNIC=brainpf0  BIN=test_multi_push bash tests/fserver/run_single_gpu.sh
kill all testing process of ps lite for user root
SCHEDULER_IP: 10.195.170.30
time: 951.45 ms
True
push_tensors: [tensor([[0.4476, 0.6482, 0.0680,  ..., 0.8657, 0.9317, 0.8116]],
       device='cuda:0')]
pull_tensors: [tensor([[0.4476, 0.6482, 0.0680,  ..., 0.8657, 0.9317, 0.8116]],
       device='cuda:0')]
worker test done
kill all testing process of ps lite for user root