Skip to content

Commit 4dbf17b

Browse files
zwischenraumAlay Shahalay2shah
authored
Fix multi-image-per-sample GRPO for LFM2-VL (#29)
* vlm_grpo: fix multi-image-per-sample for LFM2-VL TRL's GRPO buffering (split_pixel_values_by_grid -> split_tensor_dict -> unsplit_pixel_values_by_grid) does not recognise LFM2-VL's layout, where pixel_values, spatial_shapes (aliased image_sizes) and pixel_attention_mask are indexed by image rather than by sample. The stock split is a no-op for it, so split_tensor_dict slices the image axis by sample count: with >1 image/sample all but one sample's images are dropped and the logprob forward raises "Image features and image tokens do not match". LFMVLMGRPOTrainer._prepare_inputs scope-swaps an LFM2-VL-aware split/merge into TRL's namespace for the super() call (restored in finally), mirroring the existing _aliasing_spatial_shapes_as_image_sizes pattern instead of mutating TRL's module globals at import. Removable once huggingface/trl#6114 ships. Validated: 5-image dense GRPO smoke trains (0 mismatch errors). * deps: pin backend install profiles * fa2: add required attention validation * runtime: support ROCm Ray on Slurm * training: avoid eager GRPO backend imports * style: apply pre-commit formatting * Split CUDA and ROCm install profiles * Add uv backend project aliases * deps: pin TRL to 1.2.0 * runtime: clean up ROCm Ray support * deps: make FA2 install explicit * docs: clarify FA2 install flow * docs: collapse advanced install details * fix: normalize devices before FA2 env checks * test: use repo-local e2e output paths --------- Co-authored-by: Alay Shah <alay.shah@liquid.ai> Co-authored-by: alay2shah <alay0shah@gmail.com>
1 parent 0a4e9f6 commit 4dbf17b

32 files changed

Lines changed: 7260 additions & 2436 deletions

‎README.md‎

Lines changed: 113 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -47,35 +47,130 @@ git clone <repository-url>
4747
cd leap_finetune
4848
```
4949

50-
CUDA / NVIDIA clusters use the default dependency groups:
50+
### Backend Install Profiles
51+
52+
CUDA / NVIDIA clusters use the root project and default dependency groups:
53+
54+
```bash
55+
uv sync
56+
```
57+
58+
AMD / ROCm clusters use the ROCm project. Set `UV_PROJECT` once in your shell,
59+
module, or direnv config:
5160

5261
```bash
62+
export UV_PROJECT=rocm
5363
uv sync
5464
```
5565

56-
AMD / ROCm clusters should use the ROCm group instead:
66+
After `UV_PROJECT` is set, normal commands stay the same: `uv sync`,
67+
`uv run ...`, and the SLURM helpers all use the selected backend project.
5768

5869
```bash
59-
uv sync --no-group cuda --group rocm
70+
source .venv/bin/activate
71+
leap-finetune job_configs/sft_example_with_slurm.yaml
72+
73+
# ROCm, after `export UV_PROJECT=rocm`:
74+
uv run leap-finetune job_configs/sft_example_with_slurm.yaml
6075
```
6176

62-
The ROCm group is lockfile-managed and uses vLLM's ROCm wheel index for vLLM
63-
plus the matching `torch`, `torchvision`, `torchaudio`, `flash-attn`, and
64-
`triton` stack. The pinned ROCm vLLM wheels are Python 3.12 Linux wheels, so
65-
use the repo's `.python-version` when creating AMD environments.
77+
<details>
78+
<summary>Backend install details</summary>
79+
80+
The environment variable is `UV_PROJECT`, not `UV_EXPORT`.
81+
82+
The default install paths both try pinned accelerator wheels:
6683

67-
If `flash-attn` was built against a different Torch/CUDA ABI, errors such as
68-
`flash_attn_2_cuda... undefined symbol` usually mean the environment needs to
69-
be rebuilt:
84+
- CUDA: `uv sync` installs the default CUDA lock, including `vllm==0.22.0` and
85+
the pinned CUDA FlashAttention 2 wheel for the Torch 2.11 / CUDA 13 stack.
86+
- ROCm: `UV_PROJECT=rocm uv sync` installs the ROCm lock from [`rocm`](./rocm/),
87+
including direct URLs for the validated `torch`, `torchvision`, `torchaudio`,
88+
`triton`, `flash-attn`, and `vllm==0.22.0+rocm722` wheel set.
89+
90+
Ray is pinned to `2.51.1` for both profiles. The pinned accelerator wheels are
91+
Python 3.12 Linux x86_64 wheels, so use the repo's `.python-version` when
92+
creating GPU environments.
93+
94+
`UV_PROJECT` can also be used explicitly for either backend:
7095

7196
```bash
72-
uv cache clean flash-attn
73-
rm -rf .venv
74-
MAX_JOBS=1 uv sync
97+
export UV_PROJECT=cuda # optional; bare root uv is already CUDA
98+
export UV_PROJECT=rocm
99+
```
100+
101+
The top-level `cuda` path is an alias to the root project, so
102+
`UV_PROJECT=cuda` uses the same CUDA lock as bare `uv sync`.
103+
104+
No hardware-specific environment variables are required for installation. On
105+
clusters where the default uv cache is slow, quota-limited, or backed by
106+
node-local scratch, you can prefix either install command with
107+
`UV_CACHE_DIR=.uv-cache` to keep uv's package cache in the repo.
108+
109+
</details>
110+
111+
### FlashAttention 2
112+
113+
The default CUDA and ROCm install paths try to install pinned FA2 wheels. If FA2
114+
installs but does not import or cannot be selected at runtime, training emits an
115+
explicit warning and falls back to SDPA:
116+
117+
```text
118+
FlashAttention 2 not available (...); falling back to SDPA.
75119
```
76120

77-
Run this on a machine with a CUDA toolkit and enough build memory available if
78-
uv needs to rebuild `flash-attn` from source.
121+
To inspect the active environment:
122+
123+
```bash
124+
# CUDA, after `uv sync`
125+
uv run leap-finetune env fa2-status
126+
127+
# ROCm, after `export UV_PROJECT=rocm && uv sync`
128+
uv run leap-finetune env fa2-status
129+
130+
# Fail if FA2 is not usable
131+
uv run leap-finetune env fa2-status --require
132+
```
133+
134+
`fa2-status` reports the detected backend, Python tag, platform, Torch version,
135+
CUDA/HIP version, accelerator visibility, installed `flash-attn` version,
136+
selected attention implementation, and the reason FA2 is or is not usable.
137+
138+
<details>
139+
<summary>Install without FA2 or repair FA2 separately</summary>
140+
141+
If a pinned FA2 wheel cannot be resolved or installed, `uv sync` fails at install
142+
time. Install without the pinned FA2 group first, then repair FA2 separately:
143+
144+
```bash
145+
# CUDA without pinned FA2
146+
uv sync --no-group flash-attn
147+
uv run leap-finetune env install-fa2
148+
149+
# ROCm HF training without pinned FA2/vLLM
150+
UV_PROJECT=rocm uv sync --no-group rocm-fa2 --no-group rocm-vllm
151+
UV_PROJECT=rocm uv run leap-finetune env install-fa2
152+
```
153+
154+
`install-fa2` tries, in order:
155+
156+
1. A matching pinned wheel for the detected CUDA or ROCm runtime.
157+
2. Binary-only public resolution for `flash-attn==2.8.3`.
158+
3. Source build, only when explicitly requested.
159+
160+
Source builds are not part of the normal install path. Use them only as an
161+
explicit escape hatch on a machine with the matching CUDA or ROCm toolchain and
162+
enough build memory:
163+
164+
```bash
165+
uv run leap-finetune env install-fa2 --allow-source-build
166+
```
167+
168+
ROCm GRPO/vLLM support requires the full `rocm-vllm` profile and therefore a
169+
compatible ROCm FA2 stack. If that stack cannot resolve on a target cluster, use
170+
the SDPA fallback for non-vLLM training until a matching vLLM/FA2 wheel set is
171+
available.
172+
173+
</details>
79174

80175
## Quickstart
81176

@@ -828,7 +923,7 @@ Export a HuggingFace checkpoint or PEFT adapter to GGUF with
828923
`leap-export-gguf`:
829924

830925
```bash
831-
uv run leap-export-gguf /path/to/checkpoint --quant F16 --output-dir /lambdafs/gguf
926+
uv run leap-export-gguf /path/to/checkpoint --quant F16 --output-dir ./outputs/gguf
832927
```
833928

834929
Repeat `--quant` to produce multiple outputs:
@@ -837,7 +932,7 @@ Repeat `--quant` to produce multiple outputs:
837932
uv run leap-export-gguf /path/to/checkpoint \
838933
--quant F16 \
839934
--quant Q4_K_M \
840-
--output-dir /lambdafs/gguf \
935+
--output-dir ./outputs/gguf \
841936
--llama-cpp-dir /path/to/llama.cpp
842937
```
843938

@@ -852,7 +947,7 @@ PEFT adapter directories can be exported with `F16`, `BF16`, `F32`, or `Q8_0`:
852947
uv run leap-export-gguf /path/to/adapter \
853948
--base-model-path /path/to/base-model \
854949
--quant F16 \
855-
--output-dir /lambdafs/gguf
950+
--output-dir ./outputs/gguf
856951
```
857952

858953
For adapter K-quants, merge the adapter into the base model first, then export

‎cuda‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
.

‎pyproject.toml‎

Lines changed: 14 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,15 @@ dependencies = [
1111
"accelerate>=1.7.0",
1212
"peft>=0.15.2",
1313
"deepspeed>=0.17.1; sys_platform == 'linux'",
14-
"torch>=2.7.1",
14+
"torch>=2.10.0,<2.12.0",
1515
"transformers>=5.3.0,<5.4.0",
1616
"numpy>=2.4.1",
17-
"trl>=1.0,<2.0",
17+
"trl==1.2.0",
1818
"rich>=14.1.0",
1919
"pillow>=11.3.0",
2020
"mpi4py>=4.1.0; sys_platform == 'linux'",
2121
"wandb>=0.22.3",
22-
"ray==2.48.0",
22+
"ray==2.51.1",
2323
"pyyaml>=6.0",
2424
"turm>=0.13.0; sys_platform == 'linux'",
2525
"torchvision>=0.24.1",
@@ -48,72 +48,25 @@ requires = ["hatchling"]
4848
build-backend = "hatchling.build"
4949

5050
[tool.uv]
51-
# External flash-attn is optional because vLLM pins the CUDA stack tightly.
52-
# If installed via `uv sync --group flash-attn`, build against the runtime
53-
# torch ABI instead of an isolated build environment.
51+
# The default CUDA profile installs a pinned FA2 wheel. Model loading still
52+
# probes the import at runtime and falls back to SDPA if FA2 is not usable.
5453
no-build-isolation-package = ["flash-attn"]
55-
default-groups = ["dev", "cuda"]
54+
default-groups = ["dev", "cuda", "flash-attn"]
5655
environments = [
5756
"python_version == '3.12' and sys_platform == 'darwin'",
5857
"python_version == '3.12' and sys_platform == 'linux'",
5958
]
60-
# CUDA is the default GPU stack. ROCm pulls vLLM and its torch stack from a different index,
61-
# so it must be selected explicitly instead of the default cuda group:
62-
# uv sync --no-group cuda --group rocm
63-
conflicts = [
64-
[{ group = "cuda" }, { group = "rocm" }],
65-
[{ group = "flash-attn" }, { group = "rocm" }],
66-
]
6759

68-
# vLLM 0.19.0 still declares transformers<5, but LFM models require
69-
# Transformers 5.3.x. The override is intentional and covered by the GRPO
60+
# vLLM wheels can declare dependency ranges that lag the LFM Transformers /
61+
# Hugging Face Hub stack. The override is intentional and covered by the GRPO
7062
# smoke fixtures, which leave use_vllm enabled to match shipped defaults.
7163
override-dependencies = [
72-
# vllm 0.19.x declares transformers<5; override so it resolves with 5.3+
64+
# Override so vLLM resolves with the LFM-validated Transformers 5.3 stack.
7365
"transformers>=5.3.0,<5.4.0",
7466
"huggingface_hub>=1.0",
7567
"numpy<2.5",
7668
]
7769

78-
[[tool.uv.index]]
79-
name = "vllm-rocm"
80-
url = "https://wheels.vllm.ai/rocm/"
81-
explicit = true
82-
83-
[tool.uv.sources]
84-
amd-aiter = [
85-
{ index = "vllm-rocm", group = "rocm" },
86-
]
87-
amdsmi = [
88-
{ index = "vllm-rocm", group = "rocm" },
89-
]
90-
"flash-attn" = [
91-
{ index = "vllm-rocm", group = "rocm" },
92-
]
93-
torch = [
94-
{ index = "vllm-rocm", group = "rocm" },
95-
]
96-
torchaudio = [
97-
{ index = "vllm-rocm", group = "rocm" },
98-
]
99-
torchvision = [
100-
{ index = "vllm-rocm", group = "rocm" },
101-
]
102-
triton = [
103-
{ index = "vllm-rocm", group = "rocm" },
104-
]
105-
"triton-kernels" = [
106-
{ index = "vllm-rocm", group = "rocm" },
107-
]
108-
vllm = [
109-
{ index = "vllm-rocm", group = "rocm" },
110-
]
111-
112-
[[tool.uv.dependency-metadata]]
113-
name = "flash-attn"
114-
version = "2.8.3"
115-
requires-dist = ["torch", "einops", "ninja", "packaging", "setuptools", "wheel"]
116-
11770
[tool.pytest.ini_options]
11871
testpaths = ["tests"]
11972
markers = [
@@ -134,31 +87,17 @@ exclude = ["src/leap_finetune/quantization/gguf"]
13487
[dependency-groups]
13588
cuda = [
13689
"liger-kernel>=0.6.2; sys_platform == 'linux'",
137-
# Keep CUDA vLLM pinned to the Transformers 5 / LFM stack that has
138-
# been validated. External flash-attn remains opt-in below; vLLM uses
139-
# its own kernels for GRPO in the default CUDA install path.
140-
"vllm==0.19.0; sys_platform == 'linux'",
90+
# Keep CUDA vLLM pinned. FA2 is installed by default through the separate
91+
# flash-attn group so it can be disabled with `--no-group flash-attn`.
92+
"vllm==0.22.0; sys_platform == 'linux'",
14193
]
14294
dev = [
14395
"pre-commit>=4.2.0",
14496
"pytest>=8.0",
14597
"ruff>=0.11.13",
14698
]
14799
flash-attn = [
148-
# Optional HF FA2 support for environments with a known-good wheel/build.
100+
# HF FA2 support for the default CUDA 13 / Torch 2.11 stack.
149101
# Model loading still probes the import at runtime and falls back to SDPA.
150-
"flash-attn>=2.8.0; sys_platform == 'linux'",
151-
]
152-
rocm = [
153-
# Keep these direct so uv 0.6.x can apply the group-scoped vLLM ROCm index.
154-
"amd-aiter==0.1.13; sys_platform == 'linux'",
155-
"amdsmi==26.2.2; sys_platform == 'linux'",
156-
"flash-attn==2.8.3; sys_platform == 'linux'",
157-
"torch>=2.7.1; sys_platform == 'linux'",
158-
"torchaudio==2.9.0+eaa9e4e; sys_platform == 'linux'",
159-
"torchvision>=0.24.1; sys_platform == 'linux'",
160-
"triton==3.6.0; sys_platform == 'linux'",
161-
"triton-kernels==1.0.0; sys_platform == 'linux'",
162-
"liger-kernel>=0.6.2; sys_platform == 'linux'",
163-
"vllm==0.22.0+rocm722; sys_platform == 'linux'",
102+
"flash-attn @ https://github.com/adithyaxx/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu13torch2.11cxx11abiTRUE-cp312-cp312-linux_x86_64.whl#sha256=eea423825f3e12818b98b2078e2cb5ce6fe6b73d22612316d2a55fad4701938f; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
164103
]

‎rocm/pyproject.toml‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
[project]
2+
name = "leap-finetune-rocm-env"
3+
version = "0.1.0"
4+
description = "ROCm install profile for leap-finetune."
5+
requires-python = ">=3.12"
6+
dependencies = [
7+
"leap-finetune",
8+
]
9+
10+
[tool.uv]
11+
package = false
12+
no-build-isolation-package = ["flash-attn"]
13+
default-groups = ["dev", "rocm-core", "rocm-fa2", "rocm-vllm"]
14+
environments = [
15+
"python_version == '3.12' and sys_platform == 'linux'",
16+
]
17+
18+
# vLLM wheels can declare dependency ranges that lag the LFM Transformers /
19+
# Hugging Face Hub stack. Keep this aligned with the root CUDA project.
20+
override-dependencies = [
21+
"transformers>=5.3.0,<5.4.0",
22+
"huggingface_hub>=1.0",
23+
"numpy<2.5",
24+
]
25+
26+
[tool.uv.sources]
27+
leap-finetune = { path = "..", editable = true }
28+
29+
[dependency-groups]
30+
dev = [
31+
"pre-commit>=4.2.0",
32+
"pytest>=8.0",
33+
"ruff>=0.11.13",
34+
]
35+
rocm-core = [
36+
# Direct URLs freeze the validated ROCm wheel set. The vLLM ROCm simple
37+
# index is mutable and can advertise a newer torch/aiter stack for the same
38+
# package names.
39+
"torch @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/torch-2.10.0%2Bgit8514f05-cp312-cp312-manylinux_2_35_x86_64.whl#sha256=12f34bcd752decca2d7a89066d8dd6725f283fa846eeae50edab8c596880c5fc; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
40+
"torchaudio @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/torchaudio-2.9.0%2Beaa9e4e-cp312-cp312-manylinux_2_34_x86_64.whl#sha256=d15a3854246a552bdb0909aa5ee2f89cea076b0a8b17b1227d0d41d900cca1bb; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
41+
"torchvision @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/torchvision-0.24.1%2Bd801a34-cp312-cp312-manylinux_2_34_x86_64.whl#sha256=4793e928854e3372b97c03efd08e2ef3de39e83dcd322432006ea3bdb355b973; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
42+
"triton @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/triton-3.6.0-cp312-cp312-manylinux_2_35_x86_64.whl#sha256=694fd4bf8eea55c01ad8ae440fd524b6c1e34545f4539dee29eae952b5309d40; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
43+
"triton-kernels @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/triton_kernels-1.0.0-py3-none-any.whl#sha256=e986910a20025c26949f9ce2bd9985bc2fcf070b2c680e972a21bb5a0fce1032; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
44+
"liger-kernel>=0.6.2; sys_platform == 'linux'",
45+
]
46+
rocm-fa2 = [
47+
"flash-attn @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/flash_attn-2.8.3-cp312-cp312-manylinux_2_34_x86_64.whl#sha256=72bf51493106a01ac85d96493bdef3637f099c607fe2a1326f86d7b8436c89cf; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
48+
]
49+
rocm-vllm = [
50+
"amd-aiter @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/amd_aiter-0.1.13-cp312-cp312-manylinux_2_34_x86_64.whl#sha256=72c6644dc895587b69d94ae3e94012c4506aa8f261adee2f8fb034360a2c7c73; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
51+
"amdsmi @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/amdsmi-26.2.2%2B671d39a71e-py3-none-any.whl#sha256=db08928ff31c7c92c5a8555c1bfd8851575c8fdb5bc9a540f910d58ba94b6f5d; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
52+
"vllm @ https://wheels.vllm.ai/rocm/799c3afa5d5b17b676d04e0b58a5628943bb4003/vllm-0.22.0%2Brocm722-cp312-cp312-manylinux_2_34_x86_64.whl#sha256=69fec9238c0d3e0f6aae34e5a7405477eaf7ffc5fe37ed6b3edde5567d17fbb6; python_version == '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
53+
]

0 commit comments

Comments
 (0)