RFE：dockerfiles: reduce image size using multi-stage builds

[cuiyingd]

Summary
The Dockerfiles under dockerfiles install several build-time-only dependencies that remain in the final image, unnecessarily increasing its size. This impacts pull times, storage costs, and deployment efficiency.

Current Issue
The current build process installs packages that are only needed during compilation:

• build-essential: significant size impact
• ninja, packaging: unnecessary at runtime
  Other development utilities
  These packages are not required at runtime and can be removed or excluded from the final image.

Proposed Solution
Use Docker's multi-stage build feature to:
Compile dependencies (like flash-attn) in a separate builder stage with build tools installed
Copy only the compiled artifacts to a clean runtime stage
Discard build tools from the final image

Example

Stage 1: Builder

FROM nvidia/cuda:12.4 as builder
RUN apt-get install build-essential ninja packaging
RUN pip install flash-attn --no-build-isolation

Stage 2: Runtime (final image)

FROM nvidia/cuda:12.4
COPY --from=builder /usr/local/lib/python3.10/site-packages /usr/local/lib/python3.10/site-packages

No build-essential, ninja, etc. in the final image

Benefits
Smaller image size: Reduces storage and transfer time
Faster pulls: Especially important for CI/CD and multi-node clusters
Lower bandwidth: Particularly relevant for remote deployments
Best practices: Aligns with industry standards for production containers

[cuiyingd]

Let’s get the maintainer to review the RFE first to ack the requirements, and I’ll fix the issue afterward. Thanks.

[hbzfeng]

Thanks, this is a reasonable RFE.

The main caveat is that these images mix CUDA, PyTorch, conda/venv, editable installs, and RL dependencies, so the COPY site-packages approach may miss runtime pieces.

A PR would be welcome, ideally starting with the images where the benefit is clearest and including a small runtime smoke test.