alpa: OOM while serving language models

Please describe the bug Deploy Bloom-7b1 with Alpa and KubeRay. Got OOM while Bloom-7b1 is inferencing.

1:actor_name:DeviceMeshGroupManager
22022-10-26 11:54:12 | INFO | stdout | Load model alpa/bloom-7b1 ... (This can take several minutes for very large models)
32022-10-26 11:54:12 | INFO | stdout |  - Compile executables for encoder_chunk_sizes=[1, 64].
42022-10-26 11:54:19,113	WARNING worker.py:1805 -- Using blocking ray.get inside async actor. This blocks the event loop. Please use `await` on object ref with asyncio.gather if you want to yield execution to the event loop instead.
52022-10-26 11:54:47 | INFO | stdout | elapsed: 34.99 second.
62022-10-26 11:54:47 | INFO | stdout |  - Load parameters.
72022-10-26 11:55:02,213	ERROR worker.py:94 -- Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): [36mray::MeshHostWorker.load_bloom_params_worker_func()[39m (pid=501, ip=10.1.181.198, repr=<alpa.device_mesh.MeshHostWorker object at 0x7f76f5f29b20>)
8  File "/home/ray/src/llm-serving/examples/llm_serving/model/bloom_model.py", line 821, in load_bloom_params_worker_func
9    load_param(param_prefix + "self_attention.query_key_value.bias",
10  File "/home/ray/src/llm-serving/examples/llm_serving/model/bloom_model.py", line 797, in load_param
11    self.put_buffers(uuid, datas)
12  File "/home/ray/anaconda3/lib/python3.8/site-packages/alpa/device_mesh.py", line 178, in put_buffers
13    arys[batch_id][device_id] = (self.backend.buffer_from_pyval(
14jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24576 bytes.
152022-10-26 11:55:12,900	ERROR worker.py:94 -- Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): [36mray::MeshHostWorker.put_buffers()[39m (pid=501, ip=10.1.181.198, repr=<alpa.device_mesh.MeshHostWorker object at 0x7f76f5f29b20>)
16  File "/home/ray/anaconda3/lib/python3.8/site-packages/alpa/device_mesh.py", line 178, in put_buffers
17    arys[batch_id][device_id] = (self.backend.buffer_from_pyval(
18jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 33554432 bytes.
192022-10-26 11:55:22,936	ERROR worker.py:94 -- Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): [36mray::MeshHostWorker.put_buffers()[39m (pid=501, ip=10.1.181.198, repr=<alpa.device_mesh.MeshHostWorker object at 0x7f76f5f29b20>)
20  File "/home/ray/anaconda3/lib/python3.8/site-packages/alpa/device_mesh.py", line 178, in put_buffers
21    arys[batch_id][device_id] = (self.backend.buffer_from_pyval(
22jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 33554432 bytes.
23

Please describe the expected behavior

System information and environment

• OS Platform and Distribution: Linux Ubuntu 18.04.6 LTS:
• Python version: 3.8.5
• CUDA version: 11.7
• NCCL version: 2.10.3
• cupy version: cupy-cuda111==11.2.0
• GPU model and memory: NVIDIA GeForce RTX 2080 Ti, 11GB
• Alpa version: git+https://github.com/alpa-projects/alpa.git@a38bfde29e2c1ece5faf5bc59cc4189dde852091
• Ray version: 2.0.0
• KubeRay version: v0.3.0
• JAX version: jax==0.3.15, jaxlib==0.3.15+cuda111.cudnn805
• k8s version: 1.22
• 6-node cluster, with 2 GPUs on each node
• Use NFS for PVC

To Reproduce Steps to reproduce the behavior:

1. Follow this link to install kuberay operator on k8s cluster
2. Build Docker image to capture runtime environment. I put the details of the Dockerfile and build context below.
3. Prepare the YAML file for RayJob by substituting the <My Docker Image> to the image tag
4. May also need to set up the imagePullSecrets is necessary
5. Create RayJob through kubectl apply -f <RayJob yaml>. I put the content of the YAML below.
6. (Optional) Do port-forward component-alpa-service-raycluster-xxx-head-svc port 8265 to monitor the RayJob status through Ray dashboard, where xxx is some auto-generated ID.
7. Do port-forward component-alpa-service-raycluster-xxx-head-svc on port 8899
8. Send some query to the endpoint by curl -d '{"prompt":"Hello world, ","max_tokens":"128","temperature":"0.7","top_p":"0.5","model":"default"}' localhost:8899/completions

Screenshots If applicable, add screenshots to help explain your problem.

Code snippet to reproduce the problem

Additional information

• There is the YAML for RayJob

apiVersion: ray.io/v1alpha1
kind: RayJob
metadata:
  annotations:
    meta.helm.sh/release-name: component
    meta.helm.sh/release-namespace: alpa-opt-service
  creationTimestamp: "2022-10-28T15:40:49Z"
  generation: 1
  labels:
    app.kubernetes.io/managed-by: Helm
  name: component-alpa-service
  namespace: alpa-opt-service
  resourceVersion: "42668386"
  selfLink: /apis/ray.io/v1alpha1/namespaces/alpa-opt-service/rayjobs/component-alpa-service
  uid: fe665887-bbcc-4f57-9e09-3f96050e3688
spec:
  entrypoint: python start.py --model $MODEL_NAME --path $MODEL_PATH --tokenizer $TOKENIZER_NAME
    --torch-device $TORCH_DEVICE
  rayClusterSpec:
    headGroupSpec:
      rayStartParams:
        block: "true"
        dashboard-host: 0.0.0.0
        node-ip-address: $MY_POD_IP
        num-gpus: "1"
        object-store-memory: "100000000"
        port: "6379"
      replicas: 1
      serviceType: ClusterIP
      template:
        metadata:
          labels:
            app.kubernetes.io/instance: component
            app.kubernetes.io/name: kuberay
            groupName: headgroup
            rayCluster: raycluster-alpa-serving
            rayNodeType: head
        spec:
          containers:
          - env:
            - name: MY_POD_IP
              valueFrom:
                fieldRef:
                  fieldPath: status.podIP
            - name: MODEL_NAME
              value: alpa/bloom-7b1
            - name: MODEL_PATH
              value: /models
            - name: TOKENIZER_NAME
              value: bigscience/bloom-7b1
            - name: TORCH_DEVICE
              value: cpu
            image: <My Docker Image>
            imagePullPolicy: IfNotPresent
            name: ray-head
            ports:
            - containerPort: 6379
              name: gcs-server
              protocol: TCP
            - containerPort: 8265
              name: dashboard
              protocol: TCP
            - containerPort: 10001
              name: client
              protocol: TCP
            - containerPort: 8000
              name: serve
              protocol: TCP
            - containerPort: 8899
              name: alpa-service
              protocol: TCP
            resources:
              limits:
                cpu: "4"
                memory: 16Gi
                nvidia.com/gpu: "1"
              requests:
                cpu: "4"
                memory: 16Gi
          imagePullSecrets:
          - name: <My secret>
    rayVersion: 2.0.0
    workerGroupSpecs:
    - groupName: gpu
      maxReplicas: 5
      minReplicas: 2
      rayStartParams:
        block: "true"
        node-ip-address: $MY_POD_IP
        num-gpus: "1"
      replicas: 1
      template:
        spec:
          containers:
          - env:
            - name: RAY_DISABLE_DOCKER_CPU_WARNING
              value: "1"
            - name: TYPE
              value: worker
            - name: CPU_REQUEST
              valueFrom:
                resourceFieldRef:
                  containerName: machine-learning
                  resource: requests.cpu
            - name: CPU_LIMITS
              valueFrom:
                resourceFieldRef:
                  containerName: machine-learning
                  resource: limits.cpu
            - name: MEMORY_LIMITS
              valueFrom:
                resourceFieldRef:
                  containerName: machine-learning
                  resource: limits.memory
            - name: MEMORY_REQUESTS
              valueFrom:
                resourceFieldRef:
                  containerName: machine-learning
                  resource: requests.memory
            - name: MY_POD_NAME
              valueFrom:
                fieldRef:
                  fieldPath: metadata.name
            - name: MY_POD_IP
              valueFrom:
                fieldRef:
                  fieldPath: status.podIP
            image: <My Docker Image>
            imagePullPolicy: IfNotPresent
            lifecycle:
              preStop:
                exec:
                  command:
                  - /bin/sh
                  - -c
                  - ray stop
            name: machine-learning
            resources:
              limits:
                cpu: "4"
                memory: 16Gi
                nvidia.com/gpu: "1"
              requests:
                cpu: "4"
                memory: 16Gi
          imagePullSecrets:
          - name: <My secret>
          initContainers:
          - command:
            - sh
            - -c
            - until nslookup $RAY_IP.$(cat /var/run/secrets/kubernetes.io/serviceaccount/namespace).svc.cluster.local;
              do echo waiting for myservice; sleep 2; done
            image: busybox:1.28
            name: init-myservice
status:
  dashboardURL: component-alpa-service-raycluster-8d9qc-head-svc.alpa-opt-service.svc.cluster.local:8265
  endTime: "1970-01-01T00:00:00Z"
  jobDeploymentStatus: Running
  jobId: component-alpa-service-dw66f
  jobStatus: RUNNING
  message: Job is currently running.
  rayClusterName: component-alpa-service-raycluster-8d9qc
  rayClusterStatus:
    availableWorkerReplicas: 2
    desiredWorkerReplicas: 1
    endpoints:
      alpa-service: "8899"
      client: "10001"
      dashboard: "8265"
      gcs-server: "6379"
      serve: "8000"
    lastUpdateTime: "2022-10-28T15:40:54Z"
    maxWorkerReplicas: 5
    minWorkerReplicas: 2
    state: ready
  startTime: "2022-10-28T15:40:58Z"

• Details of Docker image
  • The build context contains
    • Dockerfile
    • requirements.txt
    • start.py
    • Content of Dockerfile

    FROM rayproject/ray:2.0.0-py38-cu111
    
    RUN python -m pip install --no-cache-dir --upgrade pip
    COPY requirements.txt .
    RUN python -m pip install --no-cache-dir  -r requirements.txt
    RUN python -m pip install -f https://alpa-projects.github.io/wheels.html jaxlib==0.3.15+cuda111.cudnn805
    RUN python -m pip install ray==2.0.0
    
    COPY start.py .
    

    • Content of requriements.txt

    --extra-index-url https://download.pytorch.org/whl/cu111
    torch
    cupy-cuda111
    
    git+https://github.com/alpa-projects/alpa.git@a38bfde29e2c1ece5faf5bc59cc4189dde852091 
    -e git+https://github.com/alpa-projects/alpa.git@a38bfde29e2c1ece5faf5bc59cc4189dde852091#egg=llm_serving&subdirectory=examples
    transformers==4.23.1
    fastapi
    uvicorn
    omegaconf
    jinja2
    

    • Content of start.py

import argparse
import ray
from alpa.serve import run_controller, CONTROLLER_NAME
from llm_serving.service.constants import (
    NUM_BEAMS, NUM_RETURN_SEQ, USE_RECAPTCHA)
from llm_serving.launch_model_worker import LangaugeModelWorker

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="alpa/opt-125m")
    parser.add_argument("--path", type=str, default="~/opt_weights/")
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=str, default=8899)
    parser.add_argument("--torch-device", type=str, default="cpu")
    parser.add_argument("--tokenizer", type=str)
    parser.add_argument("--no-recaptcha", action="store_true")
    parser.add_argument("--register-name", type=str, default="default")
    parser.add_argument("--ssl-keyfile", type=str)
    parser.add_argument("--ssl-certfile", type=str)
    args = parser.parse_args()

    ray.init()

    try:
        controller = ray.get_actor(CONTROLLER_NAME)
    except ValueError:
        controller = run_controller(args.host, args.port, "/",
                                    args.ssl_keyfile, args.ssl_certfile)

    group_id = 0
    controller.launch_mesh_group_manager.remote(group_id)
    t = controller.register_model.remote(
        args.register_name, LangaugeModelWorker,
        (args.model, args.path, args.torch_device, args.tokenizer, NUM_BEAMS, NUM_RETURN_SEQ,
         False if args.no_recaptcha else USE_RECAPTCHA),
        override=True)
    ray.get(t)
    t = controller.create_replica.remote(args.register_name, group_id)
    ray.get(t)

    while True:
        pass