I'm trying to run the GPT-J 6B demo available here : https://colab.research.google.com/github/kingoflolz/mesh-transformer-jax/blob/master/colab_demo.ipynb
Unfortunatelay I have some issues that are blocking me so far :
Firstly, when I'm running this part (the first cell of the colab notebook)
!apt install zstd
# the "slim" version contain only bf16 weights and no optimizer parameters, which minimizes bandwidth and memory
!time wget -c https://the-eye.eu/public/AI/GPT-J-6B/step_383500_slim.tar.zstd
!time tar -I zstd -xf step_383500_slim.tar.zstd
!git clone https://github.com/kingoflolz/mesh-transformer-jax.git
!pip install -r mesh-transformer-jax/requirements.txt
# jax 0.2.12 is required due to a regression with xmap in 0.2.13
!pip install mesh-transformer-jax/ jax==0.2.12 tensorflow==2.5.0
I don't understand why it try to download several versions of tensorflow while it's specified "tensorflow==2.5.0" in the code. Installing all these versions take a very long time. Here is a screenshot of a part of the output: output (image)
Moreover, at the end of the execution, I have this : excecution's end message (image)
Then, when trying to import the libraries in the following code cells, I receive missing modules errors. The missing modules seems to vary depending on the result of the first execution cell. missing module (image)
I believe that colab run out of disk memory trying to download model and dependencies but why this demo exists on colab if it can't be run on it ?
