Running xDiT inference engine on HAMi reports OOM

[peachest]

What happened:

When deploying task with engine thufeifeibear/xdit-dev + model black-forest-labs/FLUX.1-dev. HAMi-core reports SET_TASK_PID FAILED and Out-of-Memory error when intercepting Cuda driver API.

�[36m(pid=409)�[0m [HAMI-core ERROR (pid:409 thread=139817570217024 utils.c:160)]: host pid is error!
�[36m(pid=409)�[0m [HAMI-core Warn(409:139817570217024:libvgpu.c:857)]: SET_TASK_PID FAILED.

// Some content has been omitted......

�[36m(ImageGenerator pid=414)�[0m [HAMI-core ERROR (pid:414 thread=140470261484608 allocator.c:54)]: Device 0 OOM 48309993472 / 48305799168
�[36m(ImageGenerator pid=414)�[0m [HAMI-core ERROR (pid:414 thread=140470261484608 allocator.c:54)]: Device 0 OOM 48309993472 / 48305799168
�[36m(ImageGenerator pid=414)�[0m 
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Info(414:140470261484608:memory.c:509)]: orig free=21910454272 total=47622258688 limit=48305799168 usage=48234496000
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Info(414:140470261484608:memory.c:515)]: after free=0 total=47622258688 limit=48305799168 usage=48234496000
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:403)]: found symbol cuDeviceGetAttribute
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:403)]: found symbol cuMemAddressReserve
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:403)]: found symbol cuMemRelease
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:403)]: found symbol cuMemMap
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:403)]: found symbol cuMemCreate
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:403)]: found symbol cuMemImportFromShareableHandle
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:403)]: found symbol cuMemsetD32Async
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:libvgpu.c:74)]: into dlsym nvmlInit_v2
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:libvgpu.c:74)]: into dlsym nvmlDeviceGetHandleByPciBusId_v2
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:libvgpu.c:74)]: into dlsym nvmlDeviceGetNvLinkRemoteDeviceType
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:libvgpu.c:74)]: into dlsym nvmlDeviceGetNvLinkRemotePciInfo_v2
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:libvgpu.c:74)]: into dlsym nvmlDeviceGetComputeRunningProcesses
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:libvgpu.c:74)]: into dlsym nvmlSystemGetCudaDriverVersion_v2
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:436)]: nvmlInitWithFlags
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:hook.c:438)]: Hijacking nvmlInitWithFlags
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:device.c:10)]: Hijacking cuDeviceGetAttribute
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:device.c:10)]: Hijacking cuDeviceGetAttribute
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:device.c:10)]: Hijacking cuDeviceGetAttribute
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:device.c:10)]: Hijacking cuDeviceGetAttribute
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Debug(414:140470261484608:device.c:10)]: Hijacking cuDeviceGetAttribute
�[36m(ImageGenerator pid=414)�[0m [HAMI-core Info(414:140470261484608:hook.c:405)]: NVML DeviceGetHandleByPciBusID_v2 00000000:10:00.0
�[36m(ImageGenerator pid=414)�[0m Exception raised in creation task: The actor died because of an error raised in its creation task, �[36mray::ImageGenerator.__init__()�[39m (pid=414, ip=10.224.4.19, actor_id=4d89d46aec24713e09b28ddd01000000, repr=<server_real.ImageGenerator object at 0x7fc1c4460610>)
�[36m(ImageGenerator pid=414)�[0m            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(ImageGenerator pid=414)�[0m            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(ImageGenerator pid=414)�[0m   File "/app/server_real.py", line 71, in __init__
�[36m(ImageGenerator pid=414)�[0m     self.initialize_model(xfuser_args)
�[36m(ImageGenerator pid=414)�[0m            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(ImageGenerator pid=414)�[0m   File "/app/server_real.py", line 90, in initialize_model
�[36m(ImageGenerator pid=414)�[0m     ).to("cuda")
�[36m(ImageGenerator pid=414)�[0m       ^^^^^^^^^^
�[36m(ImageGenerator pid=414)�[0m   File "/usr/local/lib/python3.11/site-packages/xfuser/model_executor/pipelines/base_pipeline.py", line 201, in to
�[36m(ImageGenerator pid=414)�[0m     self.module = self.module.to(*args, **kwargs)
�[36m(ImageGenerator pid=414)�[0m                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(ImageGenerator pid=414)�[0m   File "/usr/local/lib/python3.11/site-packages/diffusers/pipelines/pipeline_utils.py", line 541, in to
�[36m(ImageGenerator pid=414)�[0m     module.to(device, dtype)
�[36m(ImageGenerator pid=414)�[0m   File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1369, in to
�[36m(ImageGenerator pid=414)�[0m     return self._apply(convert)
�[36m(ImageGenerator pid=414)�[0m            ^^^^^^^^^^^^^^^^^^^^
�[36m(ImageGenerator pid=414)�[0m   File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 928, in _apply
�[36m(ImageGenerator pid=414)�[0m     module._apply(fn)
�[36m(ImageGenerator pid=414)�[0m   File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 928, in _apply
�[36m(ImageGenerator pid=414)�[0m     module._apply(fn)
�[36m(ImageGenerator pid=414)�[0m   File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 928, in _apply
�[36m(ImageGenerator pid=414)�[0m     module._apply(fn)
�[36m(ImageGenerator pid=414)�[0m   [Previous line repeated 1 more time]
�[36m(ImageGenerator pid=414)�[0m   File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 955, in _apply
�[36m(ImageGenerator pid=414)�[0m     param_applied = fn(param)
�[36m(ImageGenerator pid=414)�[0m                     ^^^^^^^^^
�[36m(ImageGenerator pid=414)�[0m   File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1355, in convert
�[36m(ImageGenerator pid=414)�[0m     return t.to(
�[36m(ImageGenerator pid=414)�[0m            ^^^^^
�[36m(ImageGenerator pid=414)�[0m torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 72.00 MiB. GPU 0 has a total capacity of 44.35 GiB of which 0 bytes is free. Process 2342430 has 23.93 GiB memory in use. Of the allocated memory 23.67 GiB is allocated by PyTorch, and 6.32 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

What you expected to happen: HAMi-core working normally

How to reproduce it (as minimally and precisely as possible): deploying task with engine thufeifeibear/xdit-dev + model black-forest-labs/FLUX.1-dev

Anything else we need to know?:

Container running command:

python server_real.py --world_size 2 --ulysses_parallel_degree 2

If the Pod environment variable CUDA_DISABLE_CONTROL is set to disable HAMi-core interception, the service will operate normally without reporting an OOM error.

Set pod env variable LIBCUDA_LOG_LEVEL to 4 to print DEBUG logs of HAMi-core and saved as the following files.
Original hami-core debug log: flux.txt

To simplify reviewing bugs, the logs have been cleaned up. The results are as follows:

• Manually remove the log info of process 1 of container: flux without process 1.txt
• Manually remove log infos of all other processes except for the error process 409 of container: flux409.txt

Environment:

• HAMi version: v2.71
• nvidia driver or other AI device driver version:

When inspecting container processes via nvidia-smi upon error exit (Please view the video below.), the two worker processes occupied approximately 22–23GB each on GPU cards 2 and 3 respectively. At this point, HAMi-core's OOM error message indicated that the total GPU memory requested by a single device amounted to 45GB. It is suspected that HAMi-core aggregated the memory of both cards into a single card, triggering the OOM. It is unclear whether this is related to the initial SET_TASK_PID FAILED error.

video.mp4

[archlitchi]

Yes, i think the issue is about not counting the device memory to correct device. but i did solve a similar problem in v2.7.1, related issue: #1381, is this task launched after you updated to v2.7.1?

[peachest]

This issue has arisen in our client's production environment, where the HAMi-core they utilise was recently upgraded to version 2.7.1. We are currently trying to reproduce the problem within our internal company HAMi v2.7.1 environment and will provide further details subsequently.

[peachest]

xdit depends on the xfuser package. The flash-attn dependency within xfuer is optional. It has been verified that the service operates correctly without installing flash-attn. However, when flash-attn is installed and the service is run on HAMi, HAMi-core (v2.7.1) will report an out-of-memory (OOM) error.

Use the follow dockerfile to build v0.4.5 xdit image

# ==============================
# phase 1: building flash-attn wheel
# ==============================
FROM pytorch/pytorch:2.9.1-cuda12.8-cudnn9-devel AS builder

# Install git
RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
RUN git clone https://github.com/Dao-AILab/flash-attention.git
WORKDIR flash-attention

RUN git checkout v2.8.3 && \
    git submodule sync && \
    git submodule update --init --recursive

RUN pip install --no-cache-dir --upgrade pip setuptools wheel
RUN pip install ninja packaging psutil && ninja --version

RUN MAX_JOBS=128 pip wheel --no-cache-dir --wheel-dir=/tmp/wheels . --no-build-isolation

# ==============================
# phase 2：building xdit runtime image
# ==============================

FROM pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime AS target

# Update pip to the latest version
RUN pip install --no-cache-dir --upgrade pip

# Uninstall apex first
RUN pip uninstall -y apex

# # Install flash_attn separately with --use-pep517 flag
# RUN pip install --no-cache-dir --use-pep517 flash-attn==2.6.3 flask
RUN pip install ray==2.52.1
RUN pip install transformers==4.57.3
RUN pip install xfuser[ray]==0.4.5

WORKDIR /app

RUN pip install fastapi uvicorn
COPY ./entrypoints /app/entrypoints

EXPOSE 6000

COPY --from=builder /tmp/wheels /tmp/wheels
RUN pip install --no-cache-dir --find-links=/tmp/wheels flash-attn==2.8.3 --no-index --no-build-isolation && rm -rf /tmp/wheels
RUN pip install xfuser[flash_attn]
# Set ENTRYPOINT with CMD as default arguments
ENTRYPOINT ["python", "/app/entrypoints/launch.py"]

using the following command to build. May use --build-arg http_proxy to solve network connection problem
export the final image to k8s.io namespace

sudo nerdctl build --namespace k8s.io \
    -f docker/Dockerfile \
    -t thufeifeibear/xdit-dev:0.4.5 \
    .

---

You can run this iamage on host

sudo nerdctl run --rm -ti \
    --name xdit --network host \
    -e NVIDIA_VISIBLE_DEVICES=0,1 \
    -p 6000:6000 \
    --runtime=/usr/bin/nvidia-container-runtime \
    -v </path/to/models>.cache:/models/FLUX.1-dev \
    harbor.transwarp.io/aip/thufeifeibear/xdit-dev:0.4.5 \
        --model_path /models/FLUX.1-dev \
        --world_size 2 \
        --ulysses_parallel_degree 2 \
        --save_disk_path ""

---

Deploy on k8s using following manifests. Replace the nodeSelector and volume of hostPath as required. Makesure cpu and memory resource are enough.

---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: &app xdit
spec:
  replicas: 1
  selector:
    matchLabels:
      app: *app
  template:
    metadata:
      annotations:
      labels:
        app: *app
    spec:
      nodeSelector:
        kubernetes.io/hostname: <node with model>
      containers:
      - name: xdit-flux
        image: harbor.transwarp.io/aip/thufeifeibear/xdit-dev:0.4.5
        imagePullPolicy: IfNotPresent
        command:
        - "python"
        - "entrypoints/launch.py"
        args:
        - "--model_path=/models/FLUX.1-dev"
        - "--world_size=2"
        - "--ulysses_parallel_degree=2"
        - '--save_disk_path=""'
        resources:
          limits:
            cpu: 4
            memory: 128Gi # When ray is loading model, memory usage may spike to exceed 64GiB. Once loaded, only less than 20GiB remains resident in memory.
            nvidia.com/gpu: 2
          requests:
            cpu: 2 # With world_size set to 2, ray will launch 2 workers, thus requiring at least 2 CPUs.
            memory: 24Gi
            nvidia.com/gpu: 2
        env:
          # - name: LIBCUDA_LOG_LEVEL
          #   value: "4"    # 设置为 4，输出 HAMi-core Debug 日志。默认为 1
          # - name: CUDA_DISABLE_CONTROL
          #  value: "false" # 设置为 true，关闭 hami 拦截
          # - name: NCCL_DEBUG
          #   value: "TRACE"
          # - name: NCCL_DEBUG_SUBSYS
          #   value: "ALL"
          # - name: RAY_DEDUP_LOGS
          #   value: "0" # 设置为 0，表示关闭日志去重
        ports:
        - containerPort: 6000
        volumeMounts:
        - name: model
          mountPath: /models
        - name: shm
          mountPath: /dev/shm
      volumes:
      - name: model
        hostPath:
          path: </path/to/local/model>
      - name: shm
        emptyDir:
          medium: Memory
          sizeLimit: "32Gi"
---
apiVersion: v1
kind: Service
metadata:
  name: xdit-service
spec:
  selector:
    app: xdit
  ports:
  - protocol: TCP
    port: 6000
    targetPort: 6000
  type: NodePort
---

Here's is the nvidia-smi result screenshot