Fix IndexError with DeepSpeed ZeRO-3 when kernels rotary is active

[ArthurZucker]

Summary

Fixes #45137.

Since #41147, attention layers are decorated with @use_kernelized_func(apply_rotary_pos_emb) which attaches a rotary_fn child nn.Module at init when the kernels library is available.

DeepSpeed ZeRO-3's parameter coordinator traces the module graph at init and expects every registered submodule to run during forward. The attention forward still calls the Python apply_rotary_pos_emb, so rotary_fn is never invoked and the parameter-fetch trace desynchronizes, raising:

IndexError: pop from an empty deque
  at deepspeed/runtime/zero/partitioned_param_coordinator.py

on the second forward (reproducible via TRL's RLOO/GRPO trainers under ZeRO-3, see huggingface/trl#4899).

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[vasqu]

My only real question is: Deepspeed will no longer use a kernelized RoPE or does it?

(Also what's up with CI 👀)

[vasqu]

Unrelated? Not a big deal either way

[vasqu]

Does this mean the kernels version of RoPE cannot be used with deepspeed this way?

The whole kernelizing processes seems to go through multiple paths and it's a bit hard to keep track of:

• Kernelized RoPE only happens later (after init?)
• We register it under the module (as it's a requirement for kernels ig)

It should be fine as quick fix for sure

[ArthurZucker]

I think its the way we do functtion kernelize, (today only rope)

[ArthurZucker]

but yeah we do at init as module, only triggered lazily I think

[vasqu]

I thought a bit more about this and it might seem that this whole approach is a bit outdated? Imo, it would make more sense to use the lazy load approach as in e.g. Mamba:

transformers/src/transformers/models/mamba/modeling_mamba.py

Lines 107 to 118 in def8e6a

	global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
	causal_conv1d = lazy_load_kernel("causal-conv1d")
	causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
	causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)
	global mamba_ssm, selective_state_update, selective_scan_fn, mamba_inner_fn
	mamba_ssm = lazy_load_kernel("mamba-ssm")
	selective_state_update = resolve_internal_import(
	mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
	)
	selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None)
	mamba_inner_fn = getattr(mamba_ssm, "mamba_inner_fn", None)

No decorator anymore but keeping it in global variable state. I'm surprised either way that we had to set this as attribute to be honest 😬

[ArthurZucker]

Hhaha yeah its not outdated, that's the issue. It's "new":
This is how it would look like with your comment: #41145
vs
The PR that introduced it #41147! Reason was: #41147 (comment)