Fix blockwise sharding - Githubissues

google / jetstream-pytorch

PyTorch/XLA integration with JetStream (https://github.com/google/JetStream) for LLM inference"

Apache License 2.0

33 stars 14 forks source link

Open lsy323 opened 2 months ago

lsy323 commented 2 months ago

Fix sharding yml file for proper megatron sharding
Add weight processing hook to pad blockwise quantized weight so that the sharded dimension is divisible by the number of partitions.