[yafshar]

Change

ring_flash_attn imports flash_attn at module scope, so importing the trainer requires a flash-attn build even for runs that never touch ring attention: no-CP runs (cp=1, the default) and cp_style="ulysses" runs.

After #3533 unified context-parallel setup, the imports that eagerly load it from src/prime_rl/utils/cp.py were consolidated into one module-scope line. Move both names into the cp_style == "ring" branches that use them:

• substitute_hf_flash_attn, called only in the ring branch of setup_context_parallel
• update_ring_flash_attn_params, called only in the ring branch of setup_cp_attention_params

src/prime_rl/trainer/models/layers/ring_attn.py also imports ring_flash_attn at module scope, but substitute_ring_attn already imports that module lazily, so it is reached only on the ring path.

This matches what the surrounding code already does. flash_attn itself is guarded by try/except ImportError in src/prime_rl/trainer/models/layers/attn.py, the ulysses branches in cp.py already defer their own imports, and src/prime_rl/entrypoints/trainer.py documents deferring heavy ML imports (ring_flash_attn among them) as the intent. The ring path was the one place that still imported eagerly.

Note this is about the ring path, not the attention backend: substitute_ring_attn imports ring_flash_attn whichever flash-attn backend attn selects, so cp_style="ring" still needs it regardless of attn.

No behavior change for ring-CP runs with their dependencies installed: the same imports happen, just on first use of the ring path instead of at module import. With the dependencies missing, the error now surfaces when the ring path is first used rather than when the trainer is imported.
prime_rl._compat stays at module scope in src/prime_rl/utils/cp.py. It patches transformers for both ring_flash_attn and the mamba hub-kernels shim, and it must still run before the deferred imports.

Verification

Imported prime_rl.utils.cp before and after, on three machines with neither flash-attn nor ring-flash-attn installed (torch 2.13.0+cu130 on an L40S and a B200, torch 2.13.0+xpu on an Intel Arc Pro B60). Identical results on all three:

before: IMPORT FAILED: ModuleNotFoundError: No module named 'ring_flash_attn'
after:  IMPORT OK, leaked flash-attn modules: none
        ring path raises at call time as expected: No module named 'ring_flash_attn'

The third line is the point: the import moved rather than disappeared, so a ring-CP run still gets a clear error if required ring dependencies are missing.

Module-scope imports reachable from trainer.rl.train, trainer.sft.train and trainer.models, excluding those inside try/except ImportError:

	unguarded imports that transitively load flash-attn
before	utils/cp.py:12
after	none

The three guarded flash_attn imports in src/prime_rl/trainer/models/layers/attn.py are untouched by this PR; they already degrade to None when flash-attn is absent.

[yafshar]

@hubert-marek @mikasenghaas mind a look? Small one: the ring-CP path was the last place importing
ring_flash_attn at module scope, so a default cp=1 run needs a flash-attn build just to import
the trainer. Moved the three imports into the cp_style == "ring" branches — the ulysses branch
right below already does exactly this.

Verified on three machines with no flash-attn installed: utils/cp.py fails to import on main,
imports clean here, and the ring path still raises at call time so a real ring run gets a clear
error.

[yafshar]

Superseded by #4134, which removes ring CP and the ring-flash-attn dependency.