feat: add MoE auxiliary loss metrics
- 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:
@@ -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}")
|
||||
|
||||
@@ -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"]}
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+95
-28
@@ -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
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user