These notebooks execute on single device (CPU/GPU) via smaller dimensions to illustrate the tensor parallelism concepts covered in below medium articles.
Part-1 corresponding to tensor_parallelism_ffn_layers.ipynb - https://medium.com/@puneesh.khanna83/understanding-tensor-parallelism-to-fit-larger-models-on-multiple-devices-d4da1821d41b
Part-2 corresponding to tensor_parallelism_attention_layers.ipynb - https://medium.com/@puneesh.khanna83/understanding-tensor-parallelism-to-fit-larger-models-on-multiple-devices-part-2-ee8a2ab7f017