[DeepSeek] remove numpy, avoid tolist in gatherd_idxs

[garrett361]

Removes the numpy usage and tolist CUDA sync when computing gatherd_idxs.

[kwen2501]

Thanks much for the PR! Indeed would look much nicer!

Do you mind testing your PR with
torchrun --standalone --nproc-per-node 4 generate.py
?

I am seeing this result:

Output: You are a helpful assistant.
User: What is 2+2?

Assistant: ' ' ' ' ' ' ' ' ' ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ': ':

[garrett361]

Hmm, weird, yes I'm seeing the same. I hadn't tested generate.py, but I made some small unit tests and also checked the losses and grad norm in run.py for an randomly initialized model.

Let me look into it and I'll ping you when I figure it out.

[garrett361]

Ah, it was just the self.ep_rank * self.experts_per_rank offset I added to the gatherd_idx. The conventions here are slightly different than what I'm using elsewhere and that misunderstanding affected my unit tests, too. Removed the offset and it's fixed now:

⬢ [podman] ❯ torchrun --standalone --nproc-per-node 4 torchtitan/experiments/deepseek_v3/generate.py
W0326 19:53:05.600000 1950302 torch/distributed/run.py:792]
W0326 19:53:05.600000 1950302 torch/distributed/run.py:792] *****************************************
W0326 19:53:05.600000 1950302 torch/distributed/run.py:792] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W0326 19:53:05.600000 1950302 torch/distributed/run.py:792] *****************************************
Creating model stage 0 of 2
Creating model stage 0 of 2
Creating model stage 1 of 2
Creating model stage 1 of 2
EP rank [0]: Created Symmetric Memory for MoE
Running inference with deepseek-ai/DeepSeek-V2-Lite-Chat on (2, 2) mesh
EP rank [1]: Created Symmetric Memory for MoE
EP rank [1]: Created Symmetric Memory for MoE
EP rank [0]: Created Symmetric Memory for MoE

Output: You are a helpful assistant.
User: What is 2+2?

Assistant: 2 + 2 equals 4.

Closing inference mesh...

Apologies for that; should have tested more carefully.

[kwen2501]

Looks great! Thank you for the nice code!

[EugenHotaj]

@garrett361 Nice add! Do you notice any speed boost from removing cpu sync?

[garrett361]

Thanks @EugenHotaj ! Good question, but I didn't time it and the {run,generate}.py scripts don't have any timing metrics reported, currently, so can't say.