ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add MoE auxiliary loss metrics

Browse Source 
- Propagates MoE load-balancing loss through model outputs
- Logs task, auxiliary, and weighted losses across strategies
- Computes only explicitly requested callback metrics
- Preserves tensor compute_loss API and adds regression tests
This commit is contained in:
ViperEkura
2026-08-02 06:30:43 +08:00
parent 0fc1b1bd46
commit 1c7369f293
14 changed files with 370 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/config/train_config.py
+5 -1
View File
@@ -61,6 +61,7 @@ class TrainConfig(BaseConfig):
val_split (Optional[float]): Ratio to split from training dataset for validation, e.g. 0.05. Defaults to None.
val_step (int): Number of optimizer steps between validation runs. Defaults to 1000.
neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0.
moe_aux_loss_coef (float): Weight applied to the MoE load-balancing loss. Defaults to 0.01.
rollout_interval (int): Number of optimizer steps between online rollouts. Defaults to 512.
rollout_temperature (float): Sampling temperature for online rollout. Defaults to 0.7.
rollout_top_k (int): Top-k filtering for online rollout, 0=disable. Defaults to 0.
@@ -112,6 +113,7 @@ class TrainConfig(BaseConfig):
val_split: Optional[float] = None
val_step: int = 1000
neftune_alpha: float = 0.0
moe_aux_loss_coef: float = 0.01
rollout_interval: int = 512
rollout_temperature: float = 0.7
@@ -187,7 +189,9 @@ class TrainConfig(BaseConfig):
raise ValueError(f"rollout_top_p must be in (0, 1], got {v}")
return v
@field_validator("rollout_top_k", "num_workers", "neftune_alpha")
@field_validator(
"rollout_top_k", "num_workers", "neftune_alpha", "moe_aux_loss_coef"
)
def _validate_non_negative(cls, v):
if v < 0:
raise ValueError(f"must be non-negative, got {v}")
astrai/model/components/decoder_block.py
+11 -4
View File
@@ -1,5 +1,5 @@
from dataclasses import asdict
from typing import Optional
from typing import Optional, TypedDict
import torch.nn as nn
from torch import Tensor
@@ -10,6 +10,11 @@ from astrai.model.components.mlp import FFNFactory
from astrai.model.components.norm import RMSNorm
class DecoderOutput(TypedDict):
hidden_states: Tensor
aux_loss: Optional[Tensor]
class DecoderBlock(nn.Module):
def __init__(self, config, layer_id: int):
super().__init__()
@@ -48,7 +53,7 @@ class DecoderBlock(nn.Module):
attention_mask: Optional[Tensor] = None,
kv_cache: Optional[KVCache] = None,
is_causal: bool = False,
) -> Tensor:
) -> DecoderOutput:
attn_output = self.attention(
self.input_norm(x),
rotary_emb,
@@ -57,6 +62,8 @@ class DecoderBlock(nn.Module):
is_causal,
)
x = attn_output + x
x = self.mlp(self.post_attention_norm(x)) + x
normalized = self.post_attention_norm(x)
mlp_output = self.mlp(normalized)
x = mlp_output["hidden_states"] + x
return x
return {"hidden_states": x, "aux_loss": mlp_output["aux_loss"]}
astrai/model/components/mlp.py
+36 -11
View File
@@ -1,4 +1,4 @@
from typing import Optional
from typing import Optional, TypedDict
import torch
import torch.nn as nn
@@ -13,6 +13,16 @@ class FFNFactory(BaseFactory[nn.Module]):
pass
class FFNOutput(TypedDict):
hidden_states: Tensor
aux_loss: Optional[Tensor]
class RoutedOutput(TypedDict):
hidden_states: Tensor
aux_loss: Optional[Tensor]
@FFNFactory.register("mlp")
class MLP(nn.Module):
def __init__(self, dim: int, dim_ffn: int, down_init_std: float = 0.02):
@@ -21,10 +31,10 @@ class MLP(nn.Module):
self.gate = Linear(dim, dim_ffn)
self.down = Linear(dim_ffn, dim, init_std=down_init_std)
def forward(self, x: Tensor) -> Tensor:
def forward(self, x: Tensor) -> FFNOutput:
gated = self.up(x) * F.silu(self.gate(x))
out = self.down(gated)
return out
return {"hidden_states": out, "aux_loss": None}
@FFNFactory.register("moe")
@@ -76,22 +86,26 @@ class DeepSeekMoE(nn.Module):
]
)
def forward(self, x: Tensor) -> Tensor:
def forward(self, x: Tensor) -> FFNOutput:
include_aux_loss = self.training and torch.is_grad_enabled()
bsz, seq_len, dim = x.shape
x_flat = x.view(-1, dim)
shared_out = self._shared_forward(x_flat)
routed_out = self._routed_forward(x_flat)
routed_output = self._routed_forward(x_flat, include_aux_loss)
out = (shared_out + routed_out).view(bsz, seq_len, dim)
return out
out = (shared_out + routed_output["hidden_states"]).view(bsz, seq_len, dim)
return {"hidden_states": out, "aux_loss": routed_output["aux_loss"]}
def _shared_forward(self, x: Tensor) -> Tensor:
if self.n_shared_experts == 0:
return torch.zeros_like(x)
return sum(e(x) for e in self.shared_experts) / self.n_shared_experts
return (
sum(e(x)["hidden_states"] for e in self.shared_experts)
/ self.n_shared_experts
)
def _routed_forward(self, x: Tensor) -> Tensor:
def _routed_forward(self, x: Tensor, include_aux_loss: bool) -> RoutedOutput:
N, D = x.shape
K = self.n_activated_experts
@@ -102,15 +116,26 @@ class DeepSeekMoE(nn.Module):
if self.norm_topk_prob:
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
aux_loss = None
if include_aux_loss:
expert_load = F.one_hot(
topk_indices, num_classes=self.n_routed_experts
).float()
expert_load = expert_load.mean(dim=(0, 1))
router_prob = router_probs.float().mean(dim=0)
aux_loss = self.n_routed_experts * (expert_load * router_prob).sum()
output = torch.zeros(N, D, device=x.device, dtype=x.dtype)
for expert_idx in range(self.n_routed_experts):
expert_mask = topk_indices == expert_idx
token_idx, k_idx = expert_mask.nonzero(as_tuple=True)
if token_idx.numel() == 0:
continue
expert = self.routed_experts[expert_idx]
expert_input = x[token_idx]
expert_output = self.routed_experts[expert_idx](expert_input)
expert_output = expert(expert_input)["hidden_states"]
weights = topk_weights[token_idx, k_idx].unsqueeze(-1)
output.index_add_(0, token_idx, expert_output * weights)
return output
return {"hidden_states": output, "aux_loss": aux_loss}
astrai/model/encoder.py
+1 -1
View File
@@ -70,7 +70,7 @@ class EmbeddingEncoder(AutoModel):
attn_mask = process_attention_mask(input_mask)
for layer in self.layers:
x = layer(x, rotary_emb, attn_mask)
x = layer(x, rotary_emb, attn_mask)["hidden_states"]
hidden_states = self.norm(x)
astrai/model/transformer.py
+15 -2
View File
@@ -113,10 +113,23 @@ class AutoRegressiveLM(AutoModel):
attn_mask = process_attention_mask(input_mask)
use_sdpa_causal_mask = attn_mask is None
aux_losses = []
for layer in self.layers:
x = layer(x, rotary_emb, attn_mask, kv_cache, use_sdpa_causal_mask)
layer_output = layer(
x,
rotary_emb,
attn_mask,
kv_cache,
use_sdpa_causal_mask,
)
x = layer_output["hidden_states"]
if layer_output["aux_loss"] is not None:
aux_losses.append(layer_output["aux_loss"])
hidden_states = self.norm(x)
logits = self.lm_head(hidden_states)
return {"logits": logits, "hidden_states": hidden_states}
output = {"logits": logits, "hidden_states": hidden_states}
if aux_losses:
output["aux_loss"] = torch.stack(aux_losses).mean()
return output
astrai/trainer/strategy.py
+95 -28
View File
@@ -1,7 +1,7 @@
"""Training strategy implementations with factory pattern."""
from abc import ABC, abstractmethod
from typing import Callable, Dict, Union
from typing import Callable, Dict, Optional, TypedDict, Union
import torch
import torch.nn as nn
@@ -13,6 +13,16 @@ from astrai.parallel.executor import broadcast_state_dict
from astrai.trainer.rollout import RolloutResult
class LossOutput(TypedDict):
loss: Tensor
metrics: Dict[str, Tensor]
class LogprobsOutput(TypedDict):
logprobs: Tensor
aux_loss: Optional[Tensor]
def move_to_device(batch: Dict[str, Tensor], device: str) -> Dict[str, Tensor]:
"""Move batch tensors to specified device with non-blocking transfer."""
return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
@@ -24,7 +34,7 @@ def get_logprobs(
attn_mask: Tensor,
loss_mask: Tensor,
reduction: str,
) -> Tensor:
) -> LogprobsOutput:
"""Compute token-wise log probabilities from model outputs.
Args:
@@ -46,10 +56,11 @@ def get_logprobs(
shifted_input_ids = input_ids[:, 1:]
shifted_loss_mask = loss_mask[:, 1:]
logits = model(
outputs = model(
input_ids[:, :-1],
attn_mask[:, :, :-1, :-1] if attn_mask.dim() == 4 else attn_mask[:, :-1],
)["logits"]
)
logits = outputs["logits"]
log_probs = torch.log_softmax(logits.float(), dim=-1)
token_logprobs = torch.gather(
@@ -57,13 +68,14 @@ def get_logprobs(
).squeeze(-1)
if reduction == "mean":
return (token_logprobs * shifted_loss_mask).sum(dim=-1) / shifted_loss_mask.sum(
logprobs = (token_logprobs * shifted_loss_mask).sum(
dim=-1
).clamp(min=1.0)
) / shifted_loss_mask.sum(dim=-1).clamp(min=1.0)
elif reduction == "sum":
return (token_logprobs * shifted_loss_mask).sum(dim=-1)
logprobs = (token_logprobs * shifted_loss_mask).sum(dim=-1)
else:
return token_logprobs * shifted_loss_mask
logprobs = token_logprobs * shifted_loss_mask
return {"logprobs": logprobs, "aux_loss": outputs.get("aux_loss")}
def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
@@ -102,6 +114,7 @@ class BaseStrategy(ABC):
self.model = model
self.device = device
self.executor = kwargs.pop("executor", None)
self.moe_aux_loss_coef = kwargs.pop("moe_aux_loss_coef", 0.01)
self.extra_kwargs = kwargs
self._rollout_runner = None
@@ -117,6 +130,33 @@ class BaseStrategy(ABC):
"""
raise NotImplementedError
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
return self._normalize_output(self.compute_loss(batch))
def _loss_output(
self,
task_loss: Tensor,
metrics: Dict[str, Tensor],
aux_loss: Optional[Tensor] = None,
) -> LossOutput:
total_loss = task_loss
if aux_loss is not None:
weighted_aux_loss = self.moe_aux_loss_coef * aux_loss
total_loss = total_loss + weighted_aux_loss
metrics["moe_aux_loss"] = aux_loss
metrics["moe_aux_loss_weighted"] = weighted_aux_loss
metrics["loss"] = total_loss
return {
"loss": total_loss,
"metrics": {name: value.detach() for name, value in metrics.items()},
}
@staticmethod
def _normalize_output(output: Union[LossOutput, Tensor]) -> LossOutput:
if isinstance(output, dict):
return output
return {"loss": output, "metrics": {"loss": output.detach()}}
def supports_online(self) -> bool:
"""Whether this strategy can operate with a rollout runner.
@@ -153,17 +193,17 @@ class BaseStrategy(ABC):
if self._rollout_runner is not None:
self._rollout_runner.step()
def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
def __call__(self, batch: Dict[str, Tensor]) -> LossOutput:
"""Run offline or online forward depending on runner injection."""
if self._rollout_runner is None:
return self.compute_loss(batch)
return self.compute_loss_output(batch)
result, is_fresh = self._rollout_runner(batch)
if is_fresh:
self._on_rollout_refresh()
train_batch = self.prepare_from_rollout(result)
return self.compute_loss(train_batch)
return self.compute_loss_output(train_batch)
class StrategyFactory(BaseFactory["BaseStrategy"]):
@@ -203,9 +243,13 @@ class SEQStrategy(BaseStrategy):
self.label_smoothing = label_smoothing
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
batch = move_to_device(batch, self.device)
input_ids, target_ids = batch["input_ids"], batch["target_ids"]
logits = self.model(input_ids=input_ids)["logits"]
outputs = self.model(input_ids=input_ids)
logits = outputs["logits"]
loss = F.cross_entropy(
input=logits.flatten(0, 1).float(),
@@ -213,7 +257,7 @@ class SEQStrategy(BaseStrategy):
label_smoothing=self.label_smoothing,
)
return loss
return self._loss_output(loss, {"task_loss": loss}, outputs.get("aux_loss"))
@StrategyFactory.register("sft")
@@ -234,6 +278,9 @@ class SFTStrategy(BaseStrategy):
self.label_smoothing = label_smoothing
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
batch = move_to_device(batch, self.device)
input_ids, target_ids, position_ids, loss_mask = (
batch["input_ids"],
@@ -245,9 +292,10 @@ class SFTStrategy(BaseStrategy):
ignore_index = -100
input_mask = make_doc_boundary_mask(position_ids)
target_ids = target_ids.masked_fill(~loss_mask, ignore_index)
logits = self.model(
outputs = self.model(
input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
)["logits"]
)
logits = outputs["logits"]
loss = F.cross_entropy(
input=logits.flatten(0, 1).float(),
@@ -256,7 +304,7 @@ class SFTStrategy(BaseStrategy):
label_smoothing=self.label_smoothing,
)
return loss
return self._loss_output(loss, {"task_loss": loss}, outputs.get("aux_loss"))
@StrategyFactory.register("dpo")
@@ -282,6 +330,9 @@ class DPOStrategy(BaseStrategy):
self.reduction = reduction
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
batch = move_to_device(batch, self.device)
chosen_ids, rejected_ids = batch["chosen"], batch["rejected"]
chosen_mask, rejected_mask = batch["chosen_mask"], batch["rejected_mask"]
@@ -297,22 +348,25 @@ class DPOStrategy(BaseStrategy):
)[None, None, :, :] # [1, 1, S, S]
full_mask = key_pad & causal # [B*2, 1, S, S] — composed
log_pi = get_logprobs(
policy_output = get_logprobs(
self.model,
concat_ids,
full_mask,
concat_loss_mask,
self.reduction,
)
log_pi = policy_output["logprobs"]
aux_loss = policy_output["aux_loss"]
with torch.no_grad():
log_ref = get_logprobs(
ref_output = get_logprobs(
self.ref_model,
concat_ids,
full_mask,
concat_loss_mask,
self.reduction,
)
log_ref = ref_output["logprobs"]
log_pi_chosen = log_pi[: chosen_ids.shape[0]]
log_pi_rejected = log_pi[chosen_ids.shape[0] :]
@@ -325,7 +379,7 @@ class DPOStrategy(BaseStrategy):
ratio_diff = pi_log_ratio - ref_log_ratio
dpo_loss = -F.logsigmoid(self.beta * ratio_diff).mean()
return dpo_loss
return self._loss_output(dpo_loss, {"dpo_loss": dpo_loss}, aux_loss)
def supports_online(self) -> bool:
return True
@@ -398,6 +452,9 @@ class GRPOStrategy(BaseStrategy):
self.old_model.load_state_dict(state_dict)
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
batch = move_to_device(batch, self.device)
prompts = batch["prompts"]
responses = batch["responses"]
@@ -438,16 +495,23 @@ class GRPOStrategy(BaseStrategy):
# get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
# Response token logprobs occupy the last ``response_len`` positions
# (the first response token is predicted from the last prompt token).
token_log_probs_policy = get_logprobs(
policy_output = get_logprobs(
self.model, full_sequences, attn_mask, full_masks, "none"
)[:, prompt_len - 1 :]
)
token_log_probs_policy = policy_output["logprobs"]
aux_loss = policy_output["aux_loss"]
token_log_probs_policy = token_log_probs_policy[:, prompt_len - 1 :]
with torch.no_grad():
token_log_probs_old = get_logprobs(
old_output = get_logprobs(
self.old_model, full_sequences, attn_mask, full_masks, "none"
)[:, prompt_len - 1 :]
token_log_probs_ref = get_logprobs(
)
token_log_probs_old = old_output["logprobs"]
token_log_probs_old = token_log_probs_old[:, prompt_len - 1 :]
ref_output = get_logprobs(
self.ref_model, full_sequences, attn_mask, full_masks, "none"
)[:, prompt_len - 1 :]
)
token_log_probs_ref = ref_output["logprobs"]
token_log_probs_ref = token_log_probs_ref[:, prompt_len - 1 :]
# Reshape to [B, G, response_len]
token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
@@ -480,9 +544,12 @@ class GRPOStrategy(BaseStrategy):
kl_per_token = r - torch.log(r + eps) - 1.0
kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count
total_loss = policy_loss + kl_penalty
return total_loss
task_loss = policy_loss + kl_penalty
return self._loss_output(
task_loss,
{"policy_loss": policy_loss, "kl_loss": kl_penalty},
aux_loss,
)
def supports_online(self) -> bool:
return True
astrai/trainer/train_callback.py
+24 -7
View File
@@ -260,11 +260,28 @@ class MetricCallback(TrainCallback):
}
def _metrics(self, context: TrainContext, names):
return {
m: self._metric_funcs[m](context)
for m in names
if self._metric_funcs[m](context) is not None
}
metrics = dict(context.metrics)
for name in names:
metric_fn = self._metric_funcs.get(name)
if metric_fn is None:
continue
value = metric_fn(context)
if value is not None:
metrics[name] = value
selected = set(context.metrics) | set(names)
selected.discard("*")
result = {name: metrics[name] for name in selected if name in metrics}
if context.world_size > 1 and dist.is_initialized() and result:
metric_names = sorted(result)
values = torch.tensor(
[result[name] for name in metric_names],
dtype=torch.float32,
device=get_current_device(),
)
dist.all_reduce(values, op=dist.ReduceOp.SUM)
values /= context.world_size
result.update(zip(metric_names, values.tolist()))
return result
@only_on_rank(0)
def _append(self, event_type: str, context: TrainContext, **extra):
@@ -286,8 +303,8 @@ class MetricCallback(TrainCallback):
with torch.no_grad():
for batch in context.val_dataloader:
loss = context.strategy(batch)
total_loss += loss.item()
loss_output = context.strategy(batch)
total_loss += loss_output["loss"].item()
num_batches += 1
if context.world_size > 1 and dist.is_initialized():
astrai/trainer/train_context.py
+2
View File
@@ -38,6 +38,7 @@ class TrainContext:
epoch: int = field(default=0)
consumed_samples: int = field(default=0)
loss: float = field(default=0.0)
metrics: Dict[str, float] = field(default_factory=dict)
grad_norm: Optional[float] = field(default=None)
grad_snr_tracker: GradSNRTracker = field(default_factory=GradSNRTracker)
val_dataloader: Optional[DataLoader] = field(default=None)
@@ -221,6 +222,7 @@ class TrainContextBuilder:
obj.load_state_dict(extra[name])
strategy_kwargs = dict(cfg.extra_kwargs)
strategy_kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef)
needs_ref = cfg.strategy in (
"dpo",
astrai/trainer/trainer.py
+7 -3
View File
@@ -82,9 +82,13 @@ class Trainer:
break
with executor.accumulate(context.model):
self._call_callbacks("on_batch_begin", context)
loss = context.strategy(batch)
context.loss = loss.item()
stand_loss = loss / executor.grad_accum_steps
loss_output = context.strategy(batch)
context.loss = loss_output["loss"].item()
context.metrics = {
name: value.item()
for name, value in loss_output["metrics"].items()
}
stand_loss = loss_output["loss"] / executor.grad_accum_steps
executor.backward(stand_loss)
context.consumed_samples += (
context.config.batch_per_device * context.world_size
docs/developer/internals.md
+12 -3
View File
@@ -81,6 +81,14 @@ Where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the per-toke
Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`.
### MoE Load Balancing
MoE layers add a differentiable load-balancing term based on mean router probabilities and top-k expert assignment frequency. The training objective is:
$$ L = L_{\text{task}} + \lambda_{\text{MoE}} L_{\text{aux}} $$
`TrainConfig.moe_aux_loss_coef` controls $\lambda_{\text{MoE}}$ (default `0.01`). The unweighted and weighted auxiliary losses are logged separately.
## Training Loop Internals
Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
@@ -92,9 +100,10 @@ on_train_begin
for batch in dataloader:
on_batch_begin
with executor.accumulate(model):
loss = strategy.compute_loss(batch)
context.loss = loss.item()
stand_loss = loss / executor.grad_accum_steps
loss_output = strategy(batch)
context.loss = loss_output["loss"].item()
context.metrics = loss_output["metrics"]
stand_loss = loss_output["loss"] / executor.grad_accum_steps
executor.backward(stand_loss)
context.consumed_samples += (
context.config.batch_per_device * context.world_size
docs/guides/training.md
+6 -3
View File
@@ -54,9 +54,10 @@ on_train_begin
for batch in dataloader:
on_batch_begin
with executor.accumulate(model):
loss = strategy.compute_loss(batch)
context.loss = loss.item()
stand_loss = loss / executor.grad_accum_steps
loss_output = strategy(batch)
context.loss = loss_output["loss"].item()
context.metrics = loss_output["metrics"]
stand_loss = loss_output["loss"] / executor.grad_accum_steps
executor.backward(stand_loss)
context.consumed_samples += (
context.config.batch_per_device * context.world_size
@@ -88,6 +89,8 @@ on_train_end
Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping` (always registered; computes grad norm, clips only when `max_grad_norm` is not `None`).
Strategies return `{"loss": Tensor, "metrics": Dict[str, Tensor]}` when called by the trainer. Built-in metrics include the task-specific loss and, for MoE models, `moe_aux_loss` plus `moe_aux_loss_weighted`. Direct `compute_loss(batch)` calls continue to return a single loss tensor.
## Strategies
### SEQ (Pre-training)
tests/module/test_forward_configs.py
+45
View File
@@ -265,6 +265,51 @@ def test_moe_defaults_preserve_normalized_routing():
assert model.layers[0].mlp.norm_topk_prob is True
def test_moe_aux_loss_only_emitted_during_training():
from astrai.config.model_config import AutoRegressiveLMConfig
config = AutoRegressiveLMConfig(
**TINY_CONFIG,
ffn_type="moe",
n_routed_experts=4,
n_shared_experts=1,
n_activated_experts=2,
topk_method="greedy",
)
model = AutoRegressiveLM(config)
input_ids = torch.randint(0, config.vocab_size, (2, 8))
outputs = model(input_ids)
assert outputs["aux_loss"].ndim == 0
assert outputs["aux_loss"].requires_grad
assert torch.isfinite(outputs["aux_loss"])
with torch.no_grad():
outputs = model(input_ids)
assert "aux_loss" not in outputs
model.eval()
outputs = model(input_ids)
assert "aux_loss" not in outputs
def test_moe_component_forward_returns_ffn_output():
from astrai.model.components.mlp import DeepSeekMoE
moe = DeepSeekMoE(
dim=8,
dim_ffn=16,
n_routed_experts=4,
n_shared_experts=1,
n_activated_experts=2,
)
output = moe(torch.randn(2, 8, 8))
assert output["hidden_states"].shape == (2, 8, 8)
assert output["aux_loss"] is not None
@pytest.mark.parametrize("decoder_sparse_step", [0, -1])
def test_moe_rejects_invalid_decoder_sparse_step(decoder_sparse_step):
from pydantic import ValidationError
tests/trainer/test_loss_output.py
+106
View File
@@ -0,0 +1,106 @@
from types import SimpleNamespace
import torch
from astrai.model.transformer import AutoRegressiveLM
from astrai.trainer.strategy import BaseStrategy, SEQStrategy
from astrai.trainer.train_callback import MetricCallback
from tests.helpers import make_tiny_config
def test_seq_strategy_combines_and_reports_moe_aux_loss(device):
config = make_tiny_config(
ffn_type="moe",
n_routed_experts=4,
n_shared_experts=1,
n_activated_experts=2,
topk_method="greedy",
)
model = AutoRegressiveLM(config).to(device=device)
strategy = SEQStrategy(model, device, moe_aux_loss_coef=0.25)
batch = {
"input_ids": torch.randint(0, config.vocab_size, (2, 8), device=device),
"target_ids": torch.randint(0, config.vocab_size, (2, 8), device=device),
}
output = strategy(batch)
legacy_loss = strategy.compute_loss(batch)
assert isinstance(legacy_loss, torch.Tensor)
assert set(output["metrics"]) == {
"loss",
"task_loss",
"moe_aux_loss",
"moe_aux_loss_weighted",
}
torch.testing.assert_close(
output["loss"],
output["metrics"]["task_loss"] + output["metrics"]["moe_aux_loss_weighted"],
)
torch.testing.assert_close(
output["metrics"]["moe_aux_loss_weighted"],
0.25 * output["metrics"]["moe_aux_loss"],
)
assert output["loss"].requires_grad
assert all(not metric.requires_grad for metric in output["metrics"].values())
def test_metric_callback_includes_dynamic_strategy_metrics(tmp_path):
callback = MetricCallback(
ckpt_dir=tmp_path,
save_interval=1,
metrics=["loss", "lr"],
)
context = SimpleNamespace(
metrics={"task_loss": 2.0, "moe_aux_loss": 1.0},
loss=2.01,
optimizer=SimpleNamespace(param_groups=[{"lr": 1e-3}]),
val_loss=None,
grad_norm=None,
grad_snr_tracker=None,
world_size=1,
)
metrics = callback._metrics(context, callback.metrics)
assert metrics == {
"loss": 2.01,
"lr": 1e-3,
"task_loss": 2.0,
"moe_aux_loss": 1.0,
}
def test_metric_callback_only_computes_requested_metrics(tmp_path):
def fail_metric(context):
_ = context
raise AssertionError("unrequested metric was computed")
callback = MetricCallback(
ckpt_dir=tmp_path,
save_interval=1,
metrics=["loss"],
)
callback._metric_funcs["grad_snr"] = fail_metric
context = SimpleNamespace(
metrics={},
loss=2.0,
world_size=1,
)
metrics = callback._metrics(context, callback.metrics)
assert metrics == {"loss": 2.0}
def test_legacy_strategy_tensor_loss_is_normalized():
class LegacyStrategy(BaseStrategy):
def compute_loss(self, batch):
return torch.tensor(2.0, requires_grad=True)
strategy = LegacyStrategy(torch.nn.Linear(1, 1), "cpu")
output = strategy({})
assert output["loss"].item() == 2.0
assert output["metrics"]["loss"].item() == 2.0
tests/trainer/test_online_strategy.py
+5 -5
View File
@@ -159,21 +159,21 @@ def test_call_without_runner_falls_back_to_compute_loss_grpo(device):
"masks": torch.ones(2, 4, 6, device=device),
"rewards": torch.randn(2, 4, device=device),
}
loss = strat(batch)
loss = strat(batch)["loss"]
assert torch.isfinite(loss).item()
def test_call_with_runner_returns_finite_loss_grpo(device):
strat = _make_grpo(device)
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
assert torch.isfinite(loss).item()
def test_call_with_runner_returns_finite_loss_dpo(device):
strat = _make_dpo(device)
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
assert torch.isfinite(loss).item()
@@ -268,7 +268,7 @@ def test_step_called_when_sync_gradients_true(device):
def test_loss_is_differentiable_dpo(device):
strat = _make_dpo(device)
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
loss.backward()
has_grad = any(
p.grad is not None and p.grad.abs().sum() > 0 for p in strat.model.parameters()
@@ -279,7 +279,7 @@ def test_loss_is_differentiable_dpo(device):
def test_ref_model_not_updated_by_backward_dpo(device):
strat = _make_dpo(device)
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
loss.backward()
for p in strat.ref_model.parameters():
assert p.grad is None