# AMD ROCm MI300X Finance LLM Training & Serving
# From: rocm/pytorch:rocm6.2_ubuntu22.04_py3.10_pytorch_release_2.3.0
FROM rocm/pytorch:rocm6.2_ubuntu22.04_py3.10_pytorch_release_2.3.0
WORKDIR /workspace

RUN apt-get update && apt-get install -y --no-install-recommends git wget vim && rm -rf /var/lib/apt/lists/*

# Python ML stack — zero code changes for ROCm vs CUDA
RUN pip install --no-cache-dir     transformers>=4.45.0 accelerate peft trl datasets     huggingface-hub optimum[amd] trackio gradio vllm     ninja==1.10.2.4

# Flash Attention 2 CK backend (MI300X-native)
RUN pip install flash-attn --no-build-isolation || echo "FA2 build skipped"

# ROCm / MI300X environment tuning
ENV NCCL_MIN_NCHANNELS=112
ENV PYTORCH_TUNABLEOP_ENABLED=1
ENV TORCHINDUCTOR_MAX_AUTOTUNE=1
ENV VLLM_USE_TRITON_FLASH_ATTN=0

CMD ["python", "train.py",      "--model_name", "Qwen/Qwen2.5-7B-Instruct",      "--dataset_name", "Josephgflowers/Finance-Instruct-500k",      "--output_dir", "/workspace/output/finance-sft",      "--hub_model_id", "shah-shazid-askary/amd-finance-llm",      "--learning_rate", "1e-5", "--num_train_epochs", "3",      "--max_seq_length", "8192", "--bf16", "--push_to_hub"]