32 different nccl channel created in the decode step of DSR1 671b (IREE PJRT as backend)

[Hanjiang9831]

I only used xla for getting the hlo graph and basic graph optimizations, then I parsed the full graph to iree pjrt for compilation and runtime. I noticed that during decode, a different NCCL collective channel is created each iteration, which means a total of 33 different channels (one for prefill, 32 for decode) are created. Is this a jit cache miss problem, or should I use jax.lax.scan instead of a python loop for the decode step?

[rdyro]

The creation of multiple channels should not be a problem unless it slows down execution.

The repo unrolls layer evaluation (instead of using jax.lax.scan) to improve decode performance, but you can absolutely fork it and run experiments on top of it!