ABot-ER-9B / task_conditioned.py

Upload verified sealed ABot-ER-9B release

543bec3 verified 24 days ago

20.7 kB

"""Task-Conditioned Adaptive Training Strategy for ABot-ER dense and MoE models.
TaskConditioned keeps the pretrained projection weights shared and adds one low-rank
residual family for 2D-QA and one for Task Plan.  A task token in ``input_ids``
selects exactly one family for the whole sample.  Residual modules are injected
during model construction when restoring a TaskConditioned checkpoint.  When starting
from an ordinary backbone checkpoint they are injected only after
``from_pretrained`` has finished, avoiding Transformers 5.x ZeRO-3 missing-key
initialization corrupting their partition metadata.
"""

from __future__ import annotations

import math
import os
import weakref
from functools import wraps
from importlib import import_module
from typing import TYPE_CHECKING, Iterable, List, Sequence, Union

import torch
from torch import nn

from logging import getLogger as get_logger


if TYPE_CHECKING:
    from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer



logger = get_logger(__name__)

TaskConditioned_QA_ROUTE = 0
TaskConditioned_TASK_PLAN_ROUTE = 1


def _parse_layers(value: Union[str, Sequence[int]]) -> List[int]:
    if isinstance(value, str):
        layers = [int(item.strip()) for item in value.split(",") if item.strip()]
    else:
        layers = [int(item) for item in value]
    if not layers or len(set(layers)) != len(layers) or min(layers) < 0:
        raise ValueError(f"Invalid TaskConditioned layer selection: {value!r}")
    return sorted(layers)


class TaskConditionedLowRankResidual(nn.Module):
    """Two task-specific low-rank residuals sharing no branch parameters."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        alpha: float,
        *,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        if rank <= 0:
            raise ValueError(f"TaskConditioned rank must be positive, got {rank}.")

        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.rank = int(rank)
        self.alpha = float(alpha)
        self.scaling = self.alpha / self.rank

        factory_kwargs = {"device": device, "dtype": dtype}
        self.qa_a = nn.Linear(self.in_features, self.rank, bias=False, **factory_kwargs)
        self.qa_b = nn.Linear(self.rank, self.out_features, bias=False, **factory_kwargs)
        self.task_plan_a = nn.Linear(self.in_features, self.rank, bias=False, **factory_kwargs)
        self.task_plan_b = nn.Linear(self.rank, self.out_features, bias=False, **factory_kwargs)
        # DeepSpeed discovers ZeRO-3 leaf modules from this attribute while it
        # installs parameter-coordinator hooks.  Mark the complete residual at
        # construction time so all four low-rank projections are gathered and
        # released as one unit.  Runtime ``is_deepspeed_zero3_enabled`` checks
        # are too late/fragile here: model construction can happen before that
        # Transformers global flag is visible, leaving child Linear hooks to
        # gather the wrong parameter while Qwen's parent projection is active.
        # The attribute is inert when DeepSpeed/ZeRO-3 is not in use.
        self._z3_leaf = True
        self.reset_parameters()
        # Qwen/Transformers calls ``post_init().apply(_initialize_weights)``
        # after attention construction.  These projections already have the
        # required LoRA-style initialization (A=random, B=zero), so mark every
        # child as initialized before that traversal.  This is also critical
        # under ZeRO-3: resetting these parameters after ``zero.Init`` has
        # partitioned them corrupts DeepSpeed's parameter-shape metadata.
        for projection in (self.qa_a, self.qa_b, self.task_plan_a, self.task_plan_b):
            projection._is_hf_initialized = True
            # Transformers 5.x selects ZeRO-3 missing-key initialization from
            # flags on the parameters themselves (state_dict(keep_vars=True)).
            projection.weight._is_hf_initialized = True
        self._is_hf_initialized = True

    def reset_parameters(self) -> None:
        nn.init.kaiming_uniform_(self.qa_a.weight, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.task_plan_a.weight, a=math.sqrt(5))
        nn.init.zeros_(self.qa_b.weight)
        nn.init.zeros_(self.task_plan_b.weight)

    def _qa(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.qa_b(self.qa_a(hidden_states)) * self.scaling

    def _task_plan(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.task_plan_b(self.task_plan_a(hidden_states)) * self.scaling

    def _parameter_snapshot(self) -> dict:
        snapshot = {"z3_leaf": getattr(self, "_z3_leaf", None)}
        for name, projection in (
            ("qa_a", self.qa_a),
            ("qa_b", self.qa_b),
            ("task_plan_a", self.task_plan_a),
            ("task_plan_b", self.task_plan_b),
        ):
            parameter = projection.weight
            ds_tensor = getattr(parameter, "ds_tensor", None)
            snapshot[name] = {
                "shape": tuple(parameter.shape),
                "ds_shape": tuple(getattr(parameter, "ds_shape", ())),
                "ds_tensor_shape": tuple(ds_tensor.shape) if ds_tensor is not None else None,
                "ds_numel": getattr(parameter, "ds_numel", None),
                "ds_id": getattr(parameter, "ds_id", None),
                "ds_status": str(getattr(parameter, "ds_status", None)),
                "leaf_owner": type(getattr(parameter, "ds_z3_leaf_module", None)).__name__
                if getattr(parameter, "ds_z3_leaf_module", None) is not None
                else None,
            }
        return snapshot

    def forward(self, hidden_states: torch.Tensor, route: Union[int, torch.Tensor]) -> torch.Tensor:
        bad_parameters = [
            name
            for name, projection in (
                ("qa_a", self.qa_a),
                ("qa_b", self.qa_b),
                ("task_plan_a", self.task_plan_a),
                ("task_plan_b", self.task_plan_b),
            )
            if projection.weight.ndim != 2
        ]
        if bad_parameters:
            raise RuntimeError(
                "TaskConditioned ZeRO parameter shape corruption before routed linear: "
                f"bad={bad_parameters}, snapshot={self._parameter_snapshot()}"
            )

        # ZeRO-3 requires every data-parallel rank to traverse parameterized
        # child modules in the same order. With per-device batch size 1 the
        # route is an ``int``; conditionally executing only one branch then
        # makes QA ranks gather qa_* while Task Plan ranks gather task_plan_*,
        # corrupting the parameter coordinator's gather sequence. Always run
        # both tiny rank-r residuals and select with differentiable gates.
        # This keeps the forward/backward graph identical across ranks; the
        # inactive branch receives an exact zero gradient.
        if isinstance(route, int):
            if route not in (TaskConditioned_QA_ROUTE, TaskConditioned_TASK_PLAN_ROUTE):
                raise ValueError(f"Unknown TaskConditioned route: {route}.")
            qa_gate = hidden_states.new_tensor(float(route == TaskConditioned_QA_ROUTE))
            task_plan_gate = hidden_states.new_tensor(float(route == TaskConditioned_TASK_PLAN_ROUTE))
            return self._qa(hidden_states) * qa_gate + self._task_plan(hidden_states) * task_plan_gate

        if route.ndim != 1:
            raise ValueError(f"TaskConditioned route must be rank-1, got shape={tuple(route.shape)}.")
        if hidden_states.shape[0] != route.shape[0]:
            if route.numel() == 1:
                scalar_route = int(route.item())
                return self.forward(hidden_states, scalar_route)
            raise ValueError(
                "TaskConditioned route batch does not match projection input: "
                f"route={tuple(route.shape)}, hidden={tuple(hidden_states.shape)}."
            )

        gate_shape = (route.shape[0],) + (1,) * (hidden_states.ndim - 1)
        qa_gate = (route == TaskConditioned_QA_ROUTE).view(gate_shape).to(dtype=hidden_states.dtype)
        task_plan_gate = (route == TaskConditioned_TASK_PLAN_ROUTE).view(gate_shape).to(dtype=hidden_states.dtype)
        if not torch.all((qa_gate + task_plan_gate) == 1):
            raise ValueError("TaskConditioned route contains values other than 0 (2D-QA) and 1 (Task Plan).")
        return self._qa(hidden_states) * qa_gate + self._task_plan(hidden_states) * task_plan_gate


class TaskConditionedRoutedLinear(nn.Linear):
    """State-dict-compatible linear projection with an explicit TaskConditioned child.
    The shared ``weight`` and ``bias`` stay directly registered on the
    projection, preserving keys such as ``q_proj.weight`` when loading a base
    checkpoint.  The residual is executed from the module's normal forward
    graph instead of a forward hook.  This ordering is required by ZeRO-3's
    parameter coordinator, which cannot safely discover a new sibling module
    while another module's post-forward hooks are running.
    """

    def __init__(self, base_projection: nn.Linear, rank: int, alpha: float) -> None:
        if not isinstance(base_projection, nn.Linear):
            raise TypeError(f"TaskConditioned requires nn.Linear projections, got {type(base_projection)!r}.")

        # Do not call nn.Linear.__init__: reuse the already-created base
        # parameters so their checkpoint keys and ZeRO metadata stay intact.
        nn.Module.__init__(self)
        self.in_features = int(base_projection.in_features)
        self.out_features = int(base_projection.out_features)
        self.weight = base_projection.weight
        if base_projection.bias is None:
            self.register_parameter("bias", None)
        else:
            self.bias = base_projection.bias
        self.task_conditioned = TaskConditionedLowRankResidual(
            self.in_features,
            self.out_features,
            rank,
            alpha,
            device=base_projection.weight.device,
            dtype=base_projection.weight.dtype,
        )
        self._task_conditioned_route = None

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        route = self._task_conditioned_route
        if route is None:
            raise RuntimeError("TaskConditioned route was not configured before an adapted projection ran.")
        shared = nn.functional.linear(hidden_states, self.weight, self.bias)
        residual = self.task_conditioned(hidden_states, route)
        return shared + residual.to(dtype=shared.dtype)


def _iter_task_conditioned_projections(model: nn.Module) -> Iterable[TaskConditionedRoutedLinear]:
    for module in model.modules():
        if isinstance(module, TaskConditionedRoutedLinear):
            yield module


def _set_model_route(model: nn.Module, input_ids: torch.Tensor | None, past_key_values) -> None:
    refs = getattr(model, "_task_conditioned_projection_refs", ())
    projections = [ref() for ref in refs if ref() is not None]
    if not projections:
        return

    config = model.config
    require_route = bool(getattr(config, "task_conditioned_require_route_token", True))
    previous_route = getattr(model, "_task_conditioned_current_route", None)
    if input_ids is None:
        if previous_route is not None and past_key_values is not None:
            return
        if require_route:
            raise ValueError("TaskConditioned requires input_ids containing exactly one task route token per sample.")
        return

    qa_token_id = int(getattr(config, "task_conditioned_qa_token_id"))
    task_plan_token_id = int(getattr(config, "task_conditioned_task_plan_token_id"))
    qa_hits = input_ids.eq(qa_token_id).sum(dim=-1)
    task_plan_hits = input_ids.eq(task_plan_token_id).sum(dim=-1)
    valid = (qa_hits + task_plan_hits).eq(1)

    if not bool(valid.all()):
        # Generation decode steps normally contain only the newest token.  Keep
        # the route established by the prefill call when a cache is present.
        if previous_route is not None and past_key_values is not None and input_ids.shape[-1] <= 1:
            return
        if require_route:
            bad_rows = (~valid).nonzero(as_tuple=False).flatten().tolist()
            raise ValueError(
                "Each TaskConditioned sample must contain exactly one of the two route tokens; "
                f"invalid batch rows={bad_rows}."
            )
        return

    route_tensor = task_plan_hits.to(dtype=torch.long)
    route: Union[int, torch.Tensor]
    if route_tensor.numel() == 1:
        route = int(route_tensor.item())
    else:
        route = route_tensor
    model._task_conditioned_current_route = route
    for projection in projections:
        projection._task_conditioned_route = route


def _configure_attention_instance(attention: nn.Module, config, layer_idx: int) -> None:
    selected_layers = set(_parse_layers(getattr(config, "task_conditioned_layers")))
    if layer_idx not in selected_layers:
        return
    if not hasattr(attention, "q_proj") or not hasattr(attention, "v_proj"):
        raise TypeError(f"TaskConditioned layer {layer_idx} does not expose q_proj/v_proj.")

    already_wrapped = (
        isinstance(attention.q_proj, TaskConditionedRoutedLinear),
        isinstance(attention.v_proj, TaskConditionedRoutedLinear),
    )
    if all(already_wrapped):
        return
    if any(already_wrapped):
        raise RuntimeError(f"TaskConditioned layer {layer_idx} is only partially wrapped: {already_wrapped}.")

    rank = int(getattr(config, "task_conditioned_rank"))
    alpha = float(getattr(config, "task_conditioned_alpha"))
    attention.q_proj = TaskConditionedRoutedLinear(attention.q_proj, rank, alpha)
    attention.v_proj = TaskConditionedRoutedLinear(attention.v_proj, rank, alpha)
    attention.q_proj.task_conditioned_layer_idx = layer_idx
    attention.q_proj.task_conditioned_projection_name = "q_proj"
    attention.v_proj.task_conditioned_layer_idx = layer_idx
    attention.v_proj.task_conditioned_projection_name = "v_proj"


def _get_task_conditioned_model_classes(config):
    """Resolve the actual HF architecture, never infer it from a model path."""

    classes = {
        "qwen3_5": ("qwen3_5", "Qwen3_5Attention", "Qwen3_5ForConditionalGeneration"),
        "qwen3_5_moe": ("qwen3_5_moe", "Qwen3_5MoeAttention", "Qwen3_5MoeForConditionalGeneration"),
    }
    model_type = getattr(config, "model_type", None)
    if model_type not in classes:
        raise ValueError(f"TaskConditioned supports qwen3_5/qwen3_5_moe, got {model_type!r}.")
    family, attention_name, top_name = classes[model_type]
    modeling = import_module(f"transformers.models.{family}.modeling_{family}")
    return getattr(modeling, attention_name), getattr(modeling, top_name)


def prepare_task_conditioned_model_class(config: "PretrainedConfig") -> None:
    """Patch architecture-specific constructors before loading their weights."""

    if not bool(getattr(config, "task_conditioned_enabled", False)):
        return
    attention_cls, top_cls = _get_task_conditioned_model_classes(config)

    if not getattr(attention_cls, "_llmtuner_task_conditioned_patched", False):
        original_attention_init = attention_cls.__init__

        @wraps(original_attention_init)
        def attention_init(self, attention_config, layer_idx):
            original_attention_init(self, attention_config, layer_idx)
            if bool(getattr(attention_config, "task_conditioned_enabled", False)) and bool(
                getattr(attention_config, "task_conditioned_inject_during_init", True)
            ):
                _configure_attention_instance(self, attention_config, layer_idx)

        attention_cls.__init__ = attention_init
        attention_cls._llmtuner_task_conditioned_patched = True

    if not getattr(top_cls, "_llmtuner_task_conditioned_patched", False):
        original_top_init = top_cls.__init__
        original_top_forward = top_cls.forward

        @wraps(original_top_init)
        def top_init(self, top_config):
            original_top_init(self, top_config)
            projections = list(_iter_task_conditioned_projections(self))
            if projections:
                self._task_conditioned_projection_refs = tuple(weakref.ref(item) for item in projections)
                self._task_conditioned_current_route = None

        @wraps(original_top_forward)
        def top_forward(self, *args, **kwargs):
            input_ids = kwargs.get("input_ids")
            if input_ids is None and args:
                input_ids = args[0]
            _set_model_route(self, input_ids, kwargs.get("past_key_values"))
            return original_top_forward(self, *args, **kwargs)

        top_cls.__init__ = top_init
        top_cls.forward = top_forward
        top_cls._llmtuner_task_conditioned_patched = True


def set_task_conditioned_injection_mode(config: "PretrainedConfig", inject_during_init: bool) -> None:
    """Select the safe TaskConditioned construction lifecycle for the next model load."""

    if not bool(getattr(config, "task_conditioned_enabled", False)):
        return
    for target in (config, getattr(config, "text_config", None)):
        if target is not None:
            setattr(target, "task_conditioned_inject_during_init", bool(inject_during_init))


def inject_task_conditioned_after_pretrained_load(model: "PreTrainedModel") -> None:
    """Inject fresh TaskConditioned tensors after loading a non-TaskConditioned backbone checkpoint.
    The Transformers missing-key initialization path has already completed at
    this point.  DeepSpeed converts the newly-created ordinary parameters to
    ZeRO tensors when the training engine is initialized.
    """

    if not bool(getattr(model.config, "task_conditioned_enabled", False)):
        return
    attention_cls, _ = _get_task_conditioned_model_classes(model.config)
    for module in list(model.modules()):
        if isinstance(module, attention_cls):
            _configure_attention_instance(module, model.config, int(getattr(module, "layer_idx")))

    projections = list(_iter_task_conditioned_projections(model))
    model._task_conditioned_projection_refs = tuple(weakref.ref(item) for item in projections)
    model._task_conditioned_current_route = None
    # Checkpoints produced from this model contain TaskConditioned state and must inject
    # the modules during their next from_pretrained call.
    set_task_conditioned_injection_mode(model.config, True)
    logger.info("[task_conditioned] injected %d routed projections after backbone load.", len(projections))


def validate_task_conditioned_model(model: "PreTrainedModel") -> None:
    if not bool(getattr(model.config, "task_conditioned_enabled", False)):
        return
    expected_layers = _parse_layers(getattr(model.config, "task_conditioned_layers"))
    projections = list(_iter_task_conditioned_projections(model))
    expected_projections = sorted(
        (layer_idx, projection_name)
        for layer_idx in expected_layers
        for projection_name in ("q_proj", "v_proj")
    )
    actual_projections = sorted(
        (int(projection.task_conditioned_layer_idx), str(projection.task_conditioned_projection_name))
        for projection in projections
    )
    if actual_projections != expected_projections:
        raise RuntimeError(
            "TaskConditioned projection mismatch: "
            f"expected={expected_projections}, actual={actual_projections}."
        )

    rank = int(getattr(model.config, "task_conditioned_rank"))
    for projection in projections:
        if projection.task_conditioned.rank != rank:
            raise RuntimeError(
                "TaskConditioned rank mismatch in "
                f"layer={projection.task_conditioned_layer_idx} {projection.task_conditioned_projection_name}."
            )

    if projections and int(os.environ.get("RANK", "0")) == 0:
        first = projections[0]
        print(
            "[task_conditioned-param-audit][pre-engine] "
            f"layer={first.task_conditioned_layer_idx} projection={first.task_conditioned_projection_name} "
            f"snapshot={first.task_conditioned._parameter_snapshot()}",
            flush=True,
        )

    logger.info(
        "[task_conditioned] validated %d adapted attention layers and %d task-specific residual projections.",
        len(expected_layers),
        len(projections) * 2,
    )

	"""Task-Conditioned Adaptive Training Strategy for ABot-ER dense and MoE models.

	TaskConditioned keeps the pretrained projection weights shared and adds one low-rank
	residual family for 2D-QA and one for Task Plan. A task token in ``input_ids``
	selects exactly one family for the whole sample. Residual modules are injected
	during model construction when restoring a TaskConditioned checkpoint. When starting
	from an ordinary backbone checkpoint they are injected only after
	``from_pretrained`` has finished, avoiding Transformers 5.x ZeRO-3 missing-key
	initialization corrupting their partition metadata.
	"""

	from __future__ import annotations

	import math
	import os
	import weakref
	from functools import wraps
	from importlib import import_module
	from typing import TYPE_CHECKING, Iterable, List, Sequence, Union

	import torch
	from torch import nn

	from logging import getLogger as get_logger


	if TYPE_CHECKING:
	from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer



	logger = get_logger(__name__)

	TaskConditioned_QA_ROUTE = 0
	TaskConditioned_TASK_PLAN_ROUTE = 1


	def _parse_layers(value: Union[str, Sequence[int]]) -> List[int]:
	if isinstance(value, str):
	layers = [int(item.strip()) for item in value.split(",") if item.strip()]
	else:
	layers = [int(item) for item in value]
	if not layers or len(set(layers)) != len(layers) or min(layers) < 0:
	raise ValueError(f"Invalid TaskConditioned layer selection: {value!r}")
	return sorted(layers)


	class TaskConditionedLowRankResidual(nn.Module):
	"""Two task-specific low-rank residuals sharing no branch parameters."""

	def __init__(
	self,
	in_features: int,
	out_features: int,
	rank: int,
	alpha: float,
	*,
	device: torch.device \| None = None,
	dtype: torch.dtype \| None = None,
	) -> None:
	super().__init__()
	if rank <= 0:
	raise ValueError(f"TaskConditioned rank must be positive, got {rank}.")

	self.in_features = int(in_features)
	self.out_features = int(out_features)
	self.rank = int(rank)
	self.alpha = float(alpha)
	self.scaling = self.alpha / self.rank

	factory_kwargs = {"device": device, "dtype": dtype}
	self.qa_a = nn.Linear(self.in_features, self.rank, bias=False, **factory_kwargs)
	self.qa_b = nn.Linear(self.rank, self.out_features, bias=False, **factory_kwargs)
	self.task_plan_a = nn.Linear(self.in_features, self.rank, bias=False, **factory_kwargs)
	self.task_plan_b = nn.Linear(self.rank, self.out_features, bias=False, **factory_kwargs)
	# DeepSpeed discovers ZeRO-3 leaf modules from this attribute while it
	# installs parameter-coordinator hooks. Mark the complete residual at
	# construction time so all four low-rank projections are gathered and
	# released as one unit. Runtime ``is_deepspeed_zero3_enabled`` checks
	# are too late/fragile here: model construction can happen before that
	# Transformers global flag is visible, leaving child Linear hooks to
	# gather the wrong parameter while Qwen's parent projection is active.
	# The attribute is inert when DeepSpeed/ZeRO-3 is not in use.
	self._z3_leaf = True
	self.reset_parameters()
	# Qwen/Transformers calls ``post_init().apply(_initialize_weights)``
	# after attention construction. These projections already have the
	# required LoRA-style initialization (A=random, B=zero), so mark every
	# child as initialized before that traversal. This is also critical
	# under ZeRO-3: resetting these parameters after ``zero.Init`` has
	# partitioned them corrupts DeepSpeed's parameter-shape metadata.
	for projection in (self.qa_a, self.qa_b, self.task_plan_a, self.task_plan_b):
	projection._is_hf_initialized = True
	# Transformers 5.x selects ZeRO-3 missing-key initialization from
	# flags on the parameters themselves (state_dict(keep_vars=True)).
	projection.weight._is_hf_initialized = True
	self._is_hf_initialized = True

	def reset_parameters(self) -> None:
	nn.init.kaiming_uniform_(self.qa_a.weight, a=math.sqrt(5))
	nn.init.kaiming_uniform_(self.task_plan_a.weight, a=math.sqrt(5))
	nn.init.zeros_(self.qa_b.weight)
	nn.init.zeros_(self.task_plan_b.weight)

	def _qa(self, hidden_states: torch.Tensor) -> torch.Tensor:
	return self.qa_b(self.qa_a(hidden_states)) * self.scaling

	def _task_plan(self, hidden_states: torch.Tensor) -> torch.Tensor:
	return self.task_plan_b(self.task_plan_a(hidden_states)) * self.scaling

	def _parameter_snapshot(self) -> dict:
	snapshot = {"z3_leaf": getattr(self, "_z3_leaf", None)}
	for name, projection in (
	("qa_a", self.qa_a),
	("qa_b", self.qa_b),
	("task_plan_a", self.task_plan_a),
	("task_plan_b", self.task_plan_b),
	):
	parameter = projection.weight
	ds_tensor = getattr(parameter, "ds_tensor", None)
	snapshot[name] = {
	"shape": tuple(parameter.shape),
	"ds_shape": tuple(getattr(parameter, "ds_shape", ())),
	"ds_tensor_shape": tuple(ds_tensor.shape) if ds_tensor is not None else None,
	"ds_numel": getattr(parameter, "ds_numel", None),
	"ds_id": getattr(parameter, "ds_id", None),
	"ds_status": str(getattr(parameter, "ds_status", None)),
	"leaf_owner": type(getattr(parameter, "ds_z3_leaf_module", None)).__name__
	if getattr(parameter, "ds_z3_leaf_module", None) is not None
	else None,
	}
	return snapshot

	def forward(self, hidden_states: torch.Tensor, route: Union[int, torch.Tensor]) -> torch.Tensor:
	bad_parameters = [
	name
	for name, projection in (
	("qa_a", self.qa_a),
	("qa_b", self.qa_b),
	("task_plan_a", self.task_plan_a),
	("task_plan_b", self.task_plan_b),
	)
	if projection.weight.ndim != 2
	]
	if bad_parameters:
	raise RuntimeError(
	"TaskConditioned ZeRO parameter shape corruption before routed linear: "
	f"bad={bad_parameters}, snapshot={self._parameter_snapshot()}"
	)

	# ZeRO-3 requires every data-parallel rank to traverse parameterized
	# child modules in the same order. With per-device batch size 1 the
	# route is an ``int``; conditionally executing only one branch then
	# makes QA ranks gather qa_* while Task Plan ranks gather task_plan_*,
	# corrupting the parameter coordinator's gather sequence. Always run
	# both tiny rank-r residuals and select with differentiable gates.
	# This keeps the forward/backward graph identical across ranks; the
	# inactive branch receives an exact zero gradient.
	if isinstance(route, int):
	if route not in (TaskConditioned_QA_ROUTE, TaskConditioned_TASK_PLAN_ROUTE):
	raise ValueError(f"Unknown TaskConditioned route: {route}.")
	qa_gate = hidden_states.new_tensor(float(route == TaskConditioned_QA_ROUTE))
	task_plan_gate = hidden_states.new_tensor(float(route == TaskConditioned_TASK_PLAN_ROUTE))
	return self._qa(hidden_states) * qa_gate + self._task_plan(hidden_states) * task_plan_gate

	if route.ndim != 1:
	raise ValueError(f"TaskConditioned route must be rank-1, got shape={tuple(route.shape)}.")
	if hidden_states.shape[0] != route.shape[0]:
	if route.numel() == 1:
	scalar_route = int(route.item())
	return self.forward(hidden_states, scalar_route)
	raise ValueError(
	"TaskConditioned route batch does not match projection input: "
	f"route={tuple(route.shape)}, hidden={tuple(hidden_states.shape)}."
	)

	gate_shape = (route.shape[0],) + (1,) * (hidden_states.ndim - 1)
	qa_gate = (route == TaskConditioned_QA_ROUTE).view(gate_shape).to(dtype=hidden_states.dtype)
	task_plan_gate = (route == TaskConditioned_TASK_PLAN_ROUTE).view(gate_shape).to(dtype=hidden_states.dtype)
	if not torch.all((qa_gate + task_plan_gate) == 1):
	raise ValueError("TaskConditioned route contains values other than 0 (2D-QA) and 1 (Task Plan).")
	return self._qa(hidden_states) * qa_gate + self._task_plan(hidden_states) * task_plan_gate


	class TaskConditionedRoutedLinear(nn.Linear):
	"""State-dict-compatible linear projection with an explicit TaskConditioned child.

	The shared ``weight`` and ``bias`` stay directly registered on the
	projection, preserving keys such as ``q_proj.weight`` when loading a base
	checkpoint. The residual is executed from the module's normal forward
	graph instead of a forward hook. This ordering is required by ZeRO-3's
	parameter coordinator, which cannot safely discover a new sibling module
	while another module's post-forward hooks are running.
	"""

	def __init__(self, base_projection: nn.Linear, rank: int, alpha: float) -> None:
	if not isinstance(base_projection, nn.Linear):
	raise TypeError(f"TaskConditioned requires nn.Linear projections, got {type(base_projection)!r}.")

	# Do not call nn.Linear.__init__: reuse the already-created base
	# parameters so their checkpoint keys and ZeRO metadata stay intact.
	nn.Module.__init__(self)
	self.in_features = int(base_projection.in_features)
	self.out_features = int(base_projection.out_features)
	self.weight = base_projection.weight
	if base_projection.bias is None:
	self.register_parameter("bias", None)
	else:
	self.bias = base_projection.bias
	self.task_conditioned = TaskConditionedLowRankResidual(
	self.in_features,
	self.out_features,
	rank,
	alpha,
	device=base_projection.weight.device,
	dtype=base_projection.weight.dtype,
	)
	self._task_conditioned_route = None

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	route = self._task_conditioned_route
	if route is None:
	raise RuntimeError("TaskConditioned route was not configured before an adapted projection ran.")
	shared = nn.functional.linear(hidden_states, self.weight, self.bias)
	residual = self.task_conditioned(hidden_states, route)
	return shared + residual.to(dtype=shared.dtype)


	def _iter_task_conditioned_projections(model: nn.Module) -> Iterable[TaskConditionedRoutedLinear]:
	for module in model.modules():
	if isinstance(module, TaskConditionedRoutedLinear):
	yield module


	def _set_model_route(model: nn.Module, input_ids: torch.Tensor \| None, past_key_values) -> None:
	refs = getattr(model, "_task_conditioned_projection_refs", ())
	projections = [ref() for ref in refs if ref() is not None]
	if not projections:
	return

	config = model.config
	require_route = bool(getattr(config, "task_conditioned_require_route_token", True))
	previous_route = getattr(model, "_task_conditioned_current_route", None)
	if input_ids is None:
	if previous_route is not None and past_key_values is not None:
	return
	if require_route:
	raise ValueError("TaskConditioned requires input_ids containing exactly one task route token per sample.")
	return

	qa_token_id = int(getattr(config, "task_conditioned_qa_token_id"))
	task_plan_token_id = int(getattr(config, "task_conditioned_task_plan_token_id"))
	qa_hits = input_ids.eq(qa_token_id).sum(dim=-1)
	task_plan_hits = input_ids.eq(task_plan_token_id).sum(dim=-1)
	valid = (qa_hits + task_plan_hits).eq(1)

	if not bool(valid.all()):
	# Generation decode steps normally contain only the newest token. Keep
	# the route established by the prefill call when a cache is present.
	if previous_route is not None and past_key_values is not None and input_ids.shape[-1] <= 1:
	return
	if require_route:
	bad_rows = (~valid).nonzero(as_tuple=False).flatten().tolist()
	raise ValueError(
	"Each TaskConditioned sample must contain exactly one of the two route tokens; "
	f"invalid batch rows={bad_rows}."
	)
	return

	route_tensor = task_plan_hits.to(dtype=torch.long)
	route: Union[int, torch.Tensor]
	if route_tensor.numel() == 1:
	route = int(route_tensor.item())
	else:
	route = route_tensor
	model._task_conditioned_current_route = route
	for projection in projections:
	projection._task_conditioned_route = route


	def _configure_attention_instance(attention: nn.Module, config, layer_idx: int) -> None:
	selected_layers = set(_parse_layers(getattr(config, "task_conditioned_layers")))
	if layer_idx not in selected_layers:
	return
	if not hasattr(attention, "q_proj") or not hasattr(attention, "v_proj"):
	raise TypeError(f"TaskConditioned layer {layer_idx} does not expose q_proj/v_proj.")

	already_wrapped = (
	isinstance(attention.q_proj, TaskConditionedRoutedLinear),
	isinstance(attention.v_proj, TaskConditionedRoutedLinear),
	)
	if all(already_wrapped):
	return
	if any(already_wrapped):
	raise RuntimeError(f"TaskConditioned layer {layer_idx} is only partially wrapped: {already_wrapped}.")

	rank = int(getattr(config, "task_conditioned_rank"))
	alpha = float(getattr(config, "task_conditioned_alpha"))
	attention.q_proj = TaskConditionedRoutedLinear(attention.q_proj, rank, alpha)
	attention.v_proj = TaskConditionedRoutedLinear(attention.v_proj, rank, alpha)
	attention.q_proj.task_conditioned_layer_idx = layer_idx
	attention.q_proj.task_conditioned_projection_name = "q_proj"
	attention.v_proj.task_conditioned_layer_idx = layer_idx
	attention.v_proj.task_conditioned_projection_name = "v_proj"


	def _get_task_conditioned_model_classes(config):
	"""Resolve the actual HF architecture, never infer it from a model path."""

	classes = {
	"qwen3_5": ("qwen3_5", "Qwen3_5Attention", "Qwen3_5ForConditionalGeneration"),
	"qwen3_5_moe": ("qwen3_5_moe", "Qwen3_5MoeAttention", "Qwen3_5MoeForConditionalGeneration"),
	}
	model_type = getattr(config, "model_type", None)
	if model_type not in classes:
	raise ValueError(f"TaskConditioned supports qwen3_5/qwen3_5_moe, got {model_type!r}.")
	family, attention_name, top_name = classes[model_type]
	modeling = import_module(f"transformers.models.{family}.modeling_{family}")
	return getattr(modeling, attention_name), getattr(modeling, top_name)


	def prepare_task_conditioned_model_class(config: "PretrainedConfig") -> None:
	"""Patch architecture-specific constructors before loading their weights."""

	if not bool(getattr(config, "task_conditioned_enabled", False)):
	return
	attention_cls, top_cls = _get_task_conditioned_model_classes(config)

	if not getattr(attention_cls, "_llmtuner_task_conditioned_patched", False):
	original_attention_init = attention_cls.__init__

	@wraps(original_attention_init)
	def attention_init(self, attention_config, layer_idx):
	original_attention_init(self, attention_config, layer_idx)
	if bool(getattr(attention_config, "task_conditioned_enabled", False)) and bool(
	getattr(attention_config, "task_conditioned_inject_during_init", True)
	):
	_configure_attention_instance(self, attention_config, layer_idx)

	attention_cls.__init__ = attention_init
	attention_cls._llmtuner_task_conditioned_patched = True

	if not getattr(top_cls, "_llmtuner_task_conditioned_patched", False):
	original_top_init = top_cls.__init__
	original_top_forward = top_cls.forward

	@wraps(original_top_init)
	def top_init(self, top_config):
	original_top_init(self, top_config)
	projections = list(_iter_task_conditioned_projections(self))
	if projections:
	self._task_conditioned_projection_refs = tuple(weakref.ref(item) for item in projections)
	self._task_conditioned_current_route = None

	@wraps(original_top_forward)
	def top_forward(self, args, *kwargs):
	input_ids = kwargs.get("input_ids")
	if input_ids is None and args:
	input_ids = args[0]
	_set_model_route(self, input_ids, kwargs.get("past_key_values"))
	return original_top_forward(self, args, *kwargs)

	top_cls.__init__ = top_init
	top_cls.forward = top_forward
	top_cls._llmtuner_task_conditioned_patched = True


	def set_task_conditioned_injection_mode(config: "PretrainedConfig", inject_during_init: bool) -> None:
	"""Select the safe TaskConditioned construction lifecycle for the next model load."""

	if not bool(getattr(config, "task_conditioned_enabled", False)):
	return
	for target in (config, getattr(config, "text_config", None)):
	if target is not None:
	setattr(target, "task_conditioned_inject_during_init", bool(inject_during_init))


	def inject_task_conditioned_after_pretrained_load(model: "PreTrainedModel") -> None:
	"""Inject fresh TaskConditioned tensors after loading a non-TaskConditioned backbone checkpoint.

	The Transformers missing-key initialization path has already completed at
	this point. DeepSpeed converts the newly-created ordinary parameters to
	ZeRO tensors when the training engine is initialized.
	"""

	if not bool(getattr(model.config, "task_conditioned_enabled", False)):
	return
	attention_cls, _ = _get_task_conditioned_model_classes(model.config)
	for module in list(model.modules()):
	if isinstance(module, attention_cls):
	_configure_attention_instance(module, model.config, int(getattr(module, "layer_idx")))

	projections = list(_iter_task_conditioned_projections(model))
	model._task_conditioned_projection_refs = tuple(weakref.ref(item) for item in projections)
	model._task_conditioned_current_route = None
	# Checkpoints produced from this model contain TaskConditioned state and must inject
	# the modules during their next from_pretrained call.
	set_task_conditioned_injection_mode(model.config, True)
	logger.info("[task_conditioned] injected %d routed projections after backbone load.", len(projections))


	def validate_task_conditioned_model(model: "PreTrainedModel") -> None:
	if not bool(getattr(model.config, "task_conditioned_enabled", False)):
	return
	expected_layers = _parse_layers(getattr(model.config, "task_conditioned_layers"))
	projections = list(_iter_task_conditioned_projections(model))
	expected_projections = sorted(
	(layer_idx, projection_name)
	for layer_idx in expected_layers
	for projection_name in ("q_proj", "v_proj")
	)
	actual_projections = sorted(
	(int(projection.task_conditioned_layer_idx), str(projection.task_conditioned_projection_name))
	for projection in projections
	)
	if actual_projections != expected_projections:
	raise RuntimeError(
	"TaskConditioned projection mismatch: "
	f"expected={expected_projections}, actual={actual_projections}."
	)

	rank = int(getattr(model.config, "task_conditioned_rank"))
	for projection in projections:
	if projection.task_conditioned.rank != rank:
	raise RuntimeError(
	"TaskConditioned rank mismatch in "
	f"layer={projection.task_conditioned_layer_idx} {projection.task_conditioned_projection_name}."
	)

	if projections and int(os.environ.get("RANK", "0")) == 0:
	first = projections[0]
	print(
	"[task_conditioned-param-audit][pre-engine] "
	f"layer={first.task_conditioned_layer_idx} projection={first.task_conditioned_projection_name} "
	f"snapshot={first.task_conditioned._parameter_snapshot()}",
	flush=True,
	)

	logger.info(
	"[task_conditioned] validated %d adapted attention layers and %d task-specific residual projections.",
	len(expected_layers),
	len(projections) * 2,
	)