feat: add support for models from ibm-fms

Initially this is focused on supporting Calico architecture models, including from `gpt_megatron` format checkpoints.

Note: Even just importing / registering FMS model classes causes the process to be forked (for a warmup for PyTorch Compile). To avoid that, we can inspect the model type listed in the config.json before doing the import of FMS.

Signed-off-by: Travis Johnson <[email protected]>

Author: tjohnson31415

Dockerfile

@@ -134,7 +134,7 @@ FROM base as test-base
 
 ARG PYTHON_VERSION
 
-RUN dnf install -y make unzip python${PYTHON_VERSION} python${PYTHON_VERSION}-pip gcc openssl-devel gcc-c++ && \
+RUN dnf install -y make unzip python${PYTHON_VERSION} python${PYTHON_VERSION}-pip gcc openssl-devel gcc-c++ git && \
     dnf clean all && \
     ln -fs /usr/bin/python${PYTHON_VERSION} /usr/bin/python3 && \
     ln -s /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python && ln -s /usr/bin/pip${PYTHON_VERSION} /usr/local/bin/pip
@@ -279,7 +279,7 @@ ARG PYTHON_VERSION
 ARG SITE_PACKAGES=/opt/tgis/lib/python${PYTHON_VERSION}/site-packages
 
 # Install C++ compiler (required at runtime when PT2_COMPILE is enabled)
-RUN dnf install -y gcc-c++ && dnf clean all \
+RUN dnf install -y gcc-c++ git && dnf clean all \
     && useradd -u 2000 tgis -m -g 0
 
 SHELL ["/bin/bash", "-c"]
@@ -312,7 +312,7 @@ RUN --mount=type=bind,from=auto-gptq-cache,src=/usr/src/auto-gptq-wheel,target=/
 # Install server
 COPY proto proto
 COPY server server
-RUN cd server && make gen-server && pip install ".[accelerate, onnx-gpu, quantize]" --no-cache-dir
+RUN cd server && make gen-server && pip install ".[accelerate, ibm-fms, onnx-gpu, quantize]" --no-cache-dir
 
 # Patch codegen model changes into transformers 4.35
 RUN cp server/transformers_patch/modeling_codegen.py ${SITE_PACKAGES}/transformers/models/codegen/modeling_codegen.py

server/poetry.lock

@@ -27,6 +27,8 @@ onnxruntime = { version = "^1.16.3", optional = true }
 onnxruntime-gpu = { version = "^1.16.3", optional = true }
 onnx = { version = "^1.15.0", optional = true }
 einops = "^0.7.0"
+ibm-fms = { version = "^0.0", optional = true }
+fms-extras = {git = "https://github.com/foundation-model-stack/fms-extras", rev = "fdb1636de4261fd4102da659ab45d3fcc33fe8ef", optional = true}
 
 # Explicitly install some transitive dependencies to avoid CVEs
 jinja2 = ">=3.1.3"
@@ -39,6 +41,7 @@ cryptography = ">=42.0.2"
 [tool.poetry.extras]
 accelerate = ["accelerate"]
 bnb = ["bitsandbytes"]
+ibm-fms = ["ibm-fms", "fms-extras"]
 onnx = ["optimum", "onnxruntime", "onnx"]
 onnx-gpu = ["optimum", "onnxruntime-gpu", "onnx"]
 # These are only required if using the quantize cli command