wip. add pplx unit test

Signed-off-by: Bill Nell <[email protected]>

Author: bnellnm

examples/offline_inference/data_parallel.py

@@ -160,7 +160,7 @@ def main(model, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,
         procs.append(proc)
     exit_code = 0
     for proc in procs:
-        proc.join(timeout=300)
+        proc.join(timeout=3000)
         if proc.exitcode is None:
             print(f"Killing process {proc.pid} that "
                   f"didn't stop within 5 minutes.")

tests/kernels/moe/test_moe.py

@@ -140,9 +140,7 @@ def torch_dispatch(
 
 def torch_combine(b_out, topk_weight, topk_ids):
     num_tokens, topk = topk_ids.shape
-
     num_experts = b_out.shape[0]
-    topk = topk_ids.shape[1]
     K = b_out.shape[-1]
     out = torch.zeros((num_tokens, K), dtype=b_out.dtype, device=b_out.device)
     expert_counts = torch.zeros(num_experts, dtype=torch.int, device=b_out.device)