| """Instance-scoped dispatch for the long-context KDA training engine.""" |
|
|
| from __future__ import annotations |
|
|
| from functools import partial, wraps |
| from threading import RLock |
|
|
| from .training_engine import TrainingEngineConfig |
|
|
|
|
| _KDA_DISPATCH_LOCK = RLock() |
|
|
|
|
| def configure_kda_mixer_runtime(mixer, engine: TrainingEngineConfig): |
| """Attach the configured KDA call to one mixer without global leakage. |
| |
| FLA's ``KimiDeltaAttention.forward`` resolves ``chunk_kda`` from its module |
| global at call time. The research trainer replaced that global once per |
| process, which works for a one-model executable but leaks when a process holds |
| models with different engines. This wrapper swaps it only for this bound |
| forward and restores it in ``finally``; the custom autograd function owns its |
| backward, so the global is not needed after forward returns. |
| """ |
|
|
| if not isinstance(engine, TrainingEngineConfig): |
| raise TypeError("KDA runtime requires TrainingEngineConfig") |
| segment_size = int(engine.recurrent_segment_tokens) |
| metadata = { |
| "format": "yunmo-next-kda-training-engine-v1", |
| "installed": False, |
| "segment_size": segment_size, |
| "allow_packed_documents": bool(engine.recurrent_packed_documents), |
| "official_forward_serial_backward": bool( |
| engine.recurrent_official_forward_serial_backward |
| ), |
| } |
| if segment_size == 0: |
| mixer._yunmo_kda_training_engine = metadata |
| return mixer |
| if not ( |
| engine.recurrent_packed_documents |
| and engine.recurrent_official_forward_serial_backward |
| ): |
| raise RuntimeError("invalid serial recurrent execution configuration") |
|
|
| import fla.layers.kda as fla_kda_layer |
|
|
| from .kda_serial import serial_chunk_kda |
|
|
| engine = partial( |
| serial_chunk_kda, |
| segment_size=segment_size, |
| chunk_size=64, |
| allow_packed_documents=True, |
| official_forward_serial_backward=True, |
| ) |
| original_forward = mixer.forward |
|
|
| @wraps(original_forward) |
| def forward(*args, **kwargs): |
| with _KDA_DISPATCH_LOCK: |
| previous = fla_kda_layer.chunk_kda |
| fla_kda_layer.chunk_kda = engine |
| try: |
| return original_forward(*args, **kwargs) |
| finally: |
| fla_kda_layer.chunk_kda = previous |
|
|
| mixer.forward = forward |
| metadata = {**metadata, "installed": True, "dispatch": "instance_scoped"} |
| mixer._yunmo_kda_training_engine = metadata |
| return mixer |
|
|
