feat: support indivisible shards for TP model loading and TPlizing. (#37220)
* feat: support uneven loading and sharding
resolve merge conflicts
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
* fix: allow for empty tensor computations
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
* test: add llama1b test case
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
* due to q_proj colwise it has to be multi of 2
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
* refactor: use slice API
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
* refactor: use slice API
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
* refactor: use slice API
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
* refactor: use slice API
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
---------
Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>