ai21labs

/

Jamba-v0.1

like 1.17k

Follow

AI21 182

I've installed mamba-ssm causal-conv1d>=1.2.0 using your instruction, but fast mamba still not avaliable. What have is do wrong?
enviromenta.yml

name: jambo4
channels:
  - pytorch
  - nvidia
  - conda-forge
  - defaults
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - asttokens=2.0.5=pyhd3eb1b0_0
  - blas=1.0=mkl
  - bzip2=1.0.8=h5eee18b_5
  - ca-certificates=2024.2.2=hbcca054_0
  - certifi=2024.2.2=pyhd8ed1ab_0
  - charset-normalizer=2.0.4=pyhd3eb1b0_0
  - comm=0.2.1=py311h06a4308_0
  - cuda-cudart=12.1.105=0
  - cuda-cupti=12.1.105=0
  - cuda-libraries=12.1.0=0
  - cuda-nvrtc=12.1.105=0
  - cuda-nvtx=12.1.105=0
  - cuda-opencl=12.4.127=0
  - cuda-runtime=12.1.0=0
  - cudatoolkit-dev=11.7.0=h1de0b5d_6
  - debugpy=1.6.7=py311h6a678d5_0
  - decorator=5.1.1=pyhd3eb1b0_0
  - executing=0.8.3=pyhd3eb1b0_0
  - expat=2.5.0=h6a678d5_0
  - ffmpeg=4.3=hf484d3e_0
  - filelock=3.13.1=py311h06a4308_0
  - freetype=2.12.1=h4a9f257_0
  - gmp=6.2.1=h295c915_3
  - gmpy2=2.1.2=py311hc9b5ff0_0
  - gnutls=3.6.15=he1e5248_0
  - idna=3.4=py311h06a4308_0
  - intel-openmp=2023.1.0=hdb19cb5_46306
  - ipykernel=6.28.0=py311h06a4308_0
  - ipython=8.20.0=py311h06a4308_0
  - jedi=0.18.1=py311h06a4308_1
  - jinja2=3.1.3=py311h06a4308_0
  - jpeg=9e=h5eee18b_1
  - jupyter_client=8.6.0=py311h06a4308_0
  - jupyter_core=5.5.0=py311h06a4308_0
  - lame=3.100=h7b6447c_0
  - lcms2=2.12=h3be6417_0
  - ld_impl_linux-64=2.38=h1181459_1
  - lerc=3.0=h295c915_0
  - libcublas=12.1.0.26=0
  - libcufft=11.0.2.4=0
  - libcufile=1.9.1.3=0
  - libcurand=10.3.5.147=0
  - libcusolver=11.4.4.55=0
  - libcusparse=12.0.2.55=0
  - libdeflate=1.17=h5eee18b_1
  - libffi=3.4.4=h6a678d5_0
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libiconv=1.16=h7f8727e_2
  - libidn2=2.3.4=h5eee18b_0
  - libjpeg-turbo=2.0.0=h9bf148f_0
  - libnpp=12.0.2.50=0
  - libnvjitlink=12.1.105=0
  - libnvjpeg=12.1.1.14=0
  - libpng=1.6.39=h5eee18b_0
  - libsodium=1.0.18=h7b6447c_0
  - libstdcxx-ng=11.2.0=h1234567_1
  - libtasn1=4.19.0=h5eee18b_0
  - libtiff=4.5.1=h6a678d5_0
  - libunistring=0.9.10=h27cfd23_0
  - libuuid=1.41.5=h5eee18b_0
  - libwebp-base=1.3.2=h5eee18b_0
  - llvm-openmp=14.0.6=h9e868ea_0
  - lz4-c=1.9.4=h6a678d5_0
  - markupsafe=2.1.3=py311h5eee18b_0
  - matplotlib-inline=0.1.6=py311h06a4308_0
  - mkl=2023.1.0=h213fc3f_46344
  - mkl-service=2.4.0=py311h5eee18b_1
  - mkl_fft=1.3.8=py311h5eee18b_0
  - mkl_random=1.2.4=py311hdb19cb5_0
  - mpc=1.1.0=h10f8cd9_1
  - mpfr=4.0.2=hb69a4c5_1
  - mpmath=1.3.0=py311h06a4308_0
  - ncurses=6.4=h6a678d5_0
  - nest-asyncio=1.6.0=py311h06a4308_0
  - nettle=3.7.3=hbbd107a_1
  - networkx=3.1=py311h06a4308_0
  - numpy=1.26.4=py311h08b1b3b_0
  - numpy-base=1.26.4=py311hf175353_0
  - openh264=2.1.1=h4ff587b_0
  - openjpeg=2.4.0=h3ad879b_0
  - openssl=1.1.1w=h7f8727e_0
  - packaging=23.2=py311h06a4308_0
  - parso=0.8.3=pyhd3eb1b0_0
  - pexpect=4.8.0=pyhd3eb1b0_3
  - pillow=10.2.0=py311h5eee18b_0
  - platformdirs=3.10.0=py311h06a4308_0
  - prompt-toolkit=3.0.43=py311h06a4308_0
  - prompt_toolkit=3.0.43=hd3eb1b0_0
  - psutil=5.9.0=py311h5eee18b_0
  - ptyprocess=0.7.0=pyhd3eb1b0_2
  - pure_eval=0.2.2=pyhd3eb1b0_0
  - pygments=2.15.1=py311h06a4308_1
  - python=3.11.0=h7a1cb2a_3
  - python-dateutil=2.8.2=pyhd3eb1b0_0
  - pytorch=2.2.2=py3.11_cuda12.1_cudnn8.9.2_0
  - pytorch-cuda=12.1=ha16c6d3_5
  - pytorch-mutex=1.0=cuda
  - pyyaml=6.0.1=py311h5eee18b_0
  - pyzmq=25.1.2=py311h6a678d5_0
  - readline=8.2=h5eee18b_0
  - requests=2.31.0=py311h06a4308_1
  - setuptools=68.2.2=py311h06a4308_0
  - six=1.16.0=pyhd3eb1b0_1
  - sqlite=3.41.2=h5eee18b_0
  - stack_data=0.2.0=pyhd3eb1b0_0
  - sympy=1.12=py311h06a4308_0
  - tbb=2021.8.0=hdb19cb5_0
  - tk=8.6.12=h1ccaba5_0
  - torchaudio=2.2.2=py311_cu121
  - torchtriton=2.2.0=py311
  - torchvision=0.17.2=py311_cu121
  - tornado=6.3.3=py311h5eee18b_0
  - traitlets=5.7.1=py311h06a4308_0
  - typing_extensions=4.9.0=py311h06a4308_1
  - urllib3=2.1.0=py311h06a4308_0
  - wcwidth=0.2.5=pyhd3eb1b0_0
  - wheel=0.41.2=py311h06a4308_0
  - xz=5.4.6=h5eee18b_0
  - yaml=0.2.5=h7b6447c_0
  - zeromq=4.3.5=h6a678d5_0
  - zlib=1.2.13=h5eee18b_0
  - zstd=1.5.5=hc292b87_0
  - pip:
      - accelerate==0.29.1
      - causal-conv1d==1.2.0.post2
      - einops==0.7.0
      - fsspec==2024.3.1
      - huggingface-hub==0.22.2
      - mamba-ssm==1.2.0.post1
      - ninja==1.11.1.1
      - pandas==2.2.1
      - pip==24.0
      - pytz==2024.1
      - regex==2023.12.25
      - safetensors==0.4.2
      - tokenizers==0.15.2
      - tqdm==4.66.2
      - transformers==4.40.0.dev0
      - tzdata==2024.1
prefix: /home/alexeyv3/.conda/envs/jambo4

code

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
                                             trust_remote_code=True,
                                             )
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")

input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=216)

print(tokenizer.batch_decode(outputs))