bug(sfttrainer): SFTTrainer and ordered_pip do not reliably install flash-attn, mamba_ssm, and causal_conv1d for fms_hf_tuning_version==3.1.0

[VassilisVassiliadis]

Issue Description

SFTTrainer experiments with fms_hf_tuning_version=3.1.0 do not reliably install flash-attn, mamba_ssm, and causal_conv1d. This leads to the ray tasks that ado-sfttrainer creates to exit with an exception:

File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/ado_actuators/sfttrainer/wrapper_fms_hf_tuning/scripts/wrapper_sfttrainer.py", line 686, in <module>
  main()
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/ado_actuators/sfttrainer/wrapper_fms_hf_tuning/scripts/wrapper_sfttrainer.py", line 643, in main
  module.parse_arguments_and_execute_wrapper(
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/ado_actuators/sfttrainer/wrapper_fms_hf_tuning/tuning_versions/at_least_3_0_0_1.py", line 49, in parse_arguments_and_execute_wrapper
  return tuning.sft_trainer.train(
         ^^^^^^^^^^^^^^^^^^^^^^^^^
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/tuning/sft_trainer.py", line 351, in train
  raise e
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/tuning/sft_trainer.py", line 331, in train
  model = model_loader(
          ^^^^^^^^^^^^^
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/transformers/models/auto/auto_factory.py", line 600, in from_pretrained
  return model_class.from_pretrained(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 317, in _wrapper
  return func(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 4999, in from_pretrained
  model = cls(config, *model_args, **model_kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/transformers/models/granite/modeling_granite.py", line 480, in __init__
  super().__init__(config)
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 2228, in __init__
  self.config._attn_implementation_internal = self._check_and_adjust_attn_implementation(
                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 2765, in _check_and_adjust_attn_implementation
  attn_implementation = self.get_correct_attn_implementation(applicable_attn_implementation, is_init_check)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 2797, in get_correct_attn_implementation
  self._flash_attn_2_can_dispatch(is_init_check)
File "/tmp/ray/session_2026-07-20_03-06-53_055727_1/runtime_resources/ordered_pip/pip/4fc74475993723d02715a276975b4d6221536f41/virtualenv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 2494, in _flash_attn_2_can_dispatch
  raise ImportError(f"{preface} the package flash_attn seems to be not installed. {install_message}")
(launch_finetune_947410ce pid=3542, ip=10.48.27.185) ImportError: FlashAttention2 has been toggled on, but it cannot be used due to the following error: the package flash_attn seems to be not installed. Please refer to the documentation of https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2 to install Flash Attention 2.

However, using a wheelhouse does result in successful measurements.

To utilize a python wheelhouse:

• populate a directory that is accessible by your GPU workers with the .whl files for torch, flash_attn, mamba_ssm, causal_conv1d and their dependencies for example /data/wheelhouse
• set the env-var PIP_FIND_LINKS: "file:///data/wheelhouse/" in your ray_runtime.yaml file

In my test, I re-activated the venv that ado-sfttrainer created for one of my failed ray tasks and then I ran:

pip wheel --no-build-isolation -w /data/wheelhouse torch==2.8.0 flash_attn==2.8.3 causal_conv1d==1.5.3.post1 mamba_ssm==2.2.6.post3

Following this wheelhouse workaround, my operation was successful.

The packages flash_attn, mamba_ssm, and causal_conv1d have a build-time dependency to torch. In vanila ray there's only one way to install these packages. You must install torch in the same virtual environment of ray prior to starting your raycluster.

As a workaround, I have implemented the ordered_pip RayRuntimeEnv plugin in ado-core (docs and code) and I use it with ado-sfttrainer. However, I found that ordered_pip does not work unless also coupled with a python wheelhouse. I don't know exactly what's wrong; I suspect that it has to do with the time it takes for pip to build the above 3 wheels prior to installing them.

How to reproduce

Steps to reproduce the behavior:

1. deploy a RayCluster with the image quay.io/ado/ado:2.0.0-ofed-py311-cu121-ofed2410v1140
2. create a space for ado-sfttrainer that uses flash_attn (default is on)
3. start an operation
4. observe that the operation fails to measure properties due to the exception ImportError: FlashAttention2 has been toggled on, but it cannot be used due to the following error: the package flash_attn seems to be not installed. Please refer to the documentation of https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2 to install Flash Attention 2.

Expected behaviour

The operation is expected to work without having to use a wheelhouse.

Screenshots/Logs

Python/ado/system info

Please include the output of:

python --version
ado version
Your OS

• Python = 3.11.20
• ado = 2.0.1.dev32+ge1f8bf03.d20260715134322.dirty
• OS linux

Note: If you installed ado in editable mode (e.g., pip install -e .) or
ran uv sync, the version metadata may not be up to date. Please reinstall to
get an accurate version number:

# If you used uv sync:
uv sync --reinstall

# If you used pip install -e:
pip install -e . --force-reinstall --no-deps

Additional information

TBD

[VassilisVassiliadis]

prior to implementing ordered_pip I opened an issue in Ray about supporting packages with build-time dependencies: ray-project/ray#56418

However, Ray did not show interest in this and suggested we install torch in the base virtual environment. The issue with this approach is that we cannot try different versions of torch, and by extension fms-hf-tuning, without using a different image for the GPU workers.

[VassilisVassiliadis]

There're a few different ways to address this problem. Here are the 2 least hacky ones:

1. distribute an image for sft-trainer and pytorch 2.8 to support running experiments with fms_hf_tuning_version in ["3.0.0.1", "3.1.0"] and document that users wanting to try out different versions of fms-hf-tuning need to prepare their own ray+torch image to contain the same version of ray as their head node worker, and the same version of torch as the one their fms-hf-tuning version needs
2. document how to populate and use a python wheelhouse directory with ado-sfttrainer (see issue description)

First option will have us create N images where N is the number of python versions we want to support and those images can only be used with certain versions of fms-hf-tuning. This is the same scenario as the one we started with when we first implemented ado-sfttrainer.

The second option moves the responsibility to configuring the rayclusters onto the users and it supports any version of torch/fms-hf-tuning/python. Also, I'm fairly certain we can provide some kind of CLI utility for seeding the python wheelhouse directory.

The first option is more user-friendly, more annoying for us, and the least flexible without sacrificing functionality for the most recent release of fms-hf-tuning. The second option, is the not as user-friendly, but it is more convenient for us, and the most flexible. It also has the benefit of reducing the time it takes to setup the venvs of GPU tasks by roughly 30 minutes. This is the penalty ray pays the 1st time it sets up a venv during a oc create operation. Of course, the user does have to spend those 30 minutes the very first time they configure their RayCluster for ado-sfttrainer or whenever they wish to target a version of fms-hf-tuning that has dependencies to other versions of the packages that require compilation.

Personally, I'm leaning towards option no 2 i.e. instructions/utility-script for setting up a RayCluster for use with ado-sfttrainer. @michael-johnston @AlessandroPomponio what do you think ?