Mixtral batching support #108
Conversation
| w = topi.reshape(w, (num_experts, red, num_shards, spatial // num_shards)) | ||
| w = topi.transpose(w, (2, 0, 1, 3)) | ||
| func = te.create_prim_func([a, w]) | ||
| return func |
There was a problem hiding this comment.
I think these functions are not used for FT + multi gpu. Can you confirm?
There was a problem hiding this comment.
will double check. doesn't it depend on disco sharding?
There was a problem hiding this comment.
For FT quantization + disco, we need to use https://github.com/vinx13/mlc-llm/blob/113bd1873cb563151ed5675730be0e53560c7ab2/mlc_llm/relax_model/commons.py#L124
It hasn't been upstreamed (we should)
There was a problem hiding this comment.
I see, q4 multigpu is probably broken. This is needed for q0f16 (we are using ft kernel for moe) though
| scores, is_ascend=False, k=self.num_experts_per_tok, index_dtype="int32" | ||
| ) # (num_tokens, top_k), (num_tokens, top_k) | ||
| expert_weights = nn.emit(expert_weights / R.sum(expert_weights, axis=-1, keepdims=True)) | ||
| flattened_indices = nn.emit(relax.op.flatten(expert_indices)) |
There was a problem hiding this comment.
This flattening causes the peak VRAM footprint for the intermediate activation to be multiplied by # experts, correct? For large batch size and large # experts, this can be problematic.
Although that allows one matmul to compute all top-k experts results, I wonder if it can be beneficial to compute it sequentially (loop over top-k) to reduce VRAM usage and memory traffic. And how much perf drops if we do that.
This PR contains various fixes and workarounds to support mixtral model. Follow up tasks will be needed to clean up this.
cc @masahi @sunggg