Add Granite and Falcon H1 logit scaling support to ASFT
Extracts logit parameter resolution into `_resolve_logit_params` helper to handle model-specific scaling overrides. Adds Granite (`logits_scaling` → `1/logits_scaling`) and Falcon H1 (`lm_head_multiplier`) support alongside existing `logit_scale`/`logit_scaling` fallback chain. Updates `effective_logits` and `compute_asft_loss` to use unified resolution logic. Adds tests verifying Granite/Falcon H1 scaling in both `effective_logits` and ASFT CE
This commit is contained in:
Can
Daniel Han
parent
35cf178783
commit
61113de556
2 changed files with 169 additions and 24 deletions
|
|
@ -136,6 +136,38 @@ class TestEffectiveLogits:
|
|||
expected = 30.0 * torch.tanh(x / 30.0)
|
||||
assert torch.allclose(result, expected, atol = 1e-6)
|
||||
|
||||
def test_reads_granite_logit_scaling(self):
|
||||
"""Test Granite logit scaling override."""
|
||||
model = SimpleNamespace(
|
||||
config = SimpleNamespace(
|
||||
model_type = "granite",
|
||||
final_logit_softcapping = 0,
|
||||
logit_scale = 2.0,
|
||||
logit_scaling = 0,
|
||||
logits_scaling = 16.0,
|
||||
)
|
||||
)
|
||||
logits = torch.randn(2, 4, 8)
|
||||
result = effective_logits(logits, model)
|
||||
expected = (1.0 / 16.0) * logits.float()
|
||||
assert torch.allclose(result, expected, atol = 1e-6)
|
||||
|
||||
def test_reads_falcon_h1_logit_scaling(self):
|
||||
"""Test Falcon H1 logit scaling override."""
|
||||
model = SimpleNamespace(
|
||||
config = SimpleNamespace(
|
||||
model_type = "falcon_h1",
|
||||
final_logit_softcapping = 0,
|
||||
logit_scale = 2.0,
|
||||
logit_scaling = 0,
|
||||
lm_head_multiplier = 3.0,
|
||||
)
|
||||
)
|
||||
logits = torch.randn(2, 4, 8)
|
||||
result = effective_logits(logits, model)
|
||||
expected = 3.0 * logits.float()
|
||||
assert torch.allclose(result, expected, atol = 1e-6)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# A2) Test fast_cross_entropy_loss_per_token
|
||||
|
|
@ -443,6 +475,94 @@ class TestComputeASFTLoss:
|
|||
assert loss.dim() == 0
|
||||
assert loss.requires_grad
|
||||
|
||||
def test_sft_mode_granite_logit_scaling(self):
|
||||
"""Test Granite logit scaling in ASFT CE path."""
|
||||
|
||||
class GraniteModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.config = SimpleNamespace(
|
||||
model_type = "granite",
|
||||
final_logit_softcapping = 0,
|
||||
logit_scale = 2.0,
|
||||
logit_scaling = 0,
|
||||
logits_scaling = 8.0,
|
||||
)
|
||||
self.embedding = nn.Embedding(16, 8)
|
||||
self.linear = nn.Linear(8, 8)
|
||||
|
||||
def forward(self, input_ids = None, **kwargs):
|
||||
embeddings = self.embedding(input_ids)
|
||||
logits = self.linear(embeddings)
|
||||
return SimpleNamespace(logits = logits)
|
||||
|
||||
model = GraniteModel()
|
||||
inputs = {
|
||||
"input_ids": torch.tensor([[1, 2, 3, 4]]),
|
||||
"labels": torch.tensor([[1, 2, 3, 4]]),
|
||||
}
|
||||
captured = {}
|
||||
|
||||
def fake_ce(logits, labels, logit_softcapping = 0, logit_scaling = 0, ignore_index = -100):
|
||||
captured["logit_scaling"] = logit_scaling
|
||||
batch, seq_len, _ = logits.shape
|
||||
losses = torch.zeros(batch * seq_len, device = logits.device)
|
||||
valid_mask = labels.view(-1) != ignore_index
|
||||
return losses, valid_mask
|
||||
|
||||
with patch(
|
||||
"unsloth.losses.asft.fast_cross_entropy_loss_per_token",
|
||||
side_effect = fake_ce,
|
||||
):
|
||||
loss = compute_asft_loss(model, inputs, asft_mode = "sft", kl_weight = 0.0)
|
||||
|
||||
assert captured["logit_scaling"] == pytest.approx(1.0 / 8.0)
|
||||
assert loss.dim() == 0
|
||||
|
||||
def test_sft_mode_falcon_h1_logit_scaling(self):
|
||||
"""Test Falcon H1 logit scaling in ASFT CE path."""
|
||||
|
||||
class FalconH1Model(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.config = SimpleNamespace(
|
||||
model_type = "falcon_h1",
|
||||
final_logit_softcapping = 0,
|
||||
logit_scale = 0,
|
||||
logit_scaling = 0,
|
||||
lm_head_multiplier = 3.0,
|
||||
)
|
||||
self.embedding = nn.Embedding(16, 8)
|
||||
self.linear = nn.Linear(8, 8)
|
||||
|
||||
def forward(self, input_ids = None, **kwargs):
|
||||
embeddings = self.embedding(input_ids)
|
||||
logits = self.linear(embeddings)
|
||||
return SimpleNamespace(logits = logits)
|
||||
|
||||
model = FalconH1Model()
|
||||
inputs = {
|
||||
"input_ids": torch.tensor([[1, 2, 3, 4]]),
|
||||
"labels": torch.tensor([[1, 2, 3, 4]]),
|
||||
}
|
||||
captured = {}
|
||||
|
||||
def fake_ce(logits, labels, logit_softcapping = 0, logit_scaling = 0, ignore_index = -100):
|
||||
captured["logit_scaling"] = logit_scaling
|
||||
batch, seq_len, _ = logits.shape
|
||||
losses = torch.zeros(batch * seq_len, device = logits.device)
|
||||
valid_mask = labels.view(-1) != ignore_index
|
||||
return losses, valid_mask
|
||||
|
||||
with patch(
|
||||
"unsloth.losses.asft.fast_cross_entropy_loss_per_token",
|
||||
side_effect = fake_ce,
|
||||
):
|
||||
loss = compute_asft_loss(model, inputs, asft_mode = "sft", kl_weight = 0.0)
|
||||
|
||||
assert captured["logit_scaling"] == pytest.approx(3.0)
|
||||
assert loss.dim() == 0
|
||||
|
||||
def test_dft_mode(self, simple_model):
|
||||
"""Test DFT mode."""
|
||||
inputs = {
|
||||
|
|
|
|||
|
|
@ -99,6 +99,45 @@ class ASFTStreamingConfig:
|
|||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _resolve_logit_params(
|
||||
model: Optional[nn.Module],
|
||||
logit_softcapping: Optional[float],
|
||||
logit_scaling: Optional[float],
|
||||
) -> Tuple[float, float]:
|
||||
if model is not None:
|
||||
config = getattr(model, "config", None)
|
||||
if config is not None:
|
||||
if logit_softcapping is None:
|
||||
logit_softcapping = getattr(config, "final_logit_softcapping", 0)
|
||||
if logit_softcapping is None:
|
||||
logit_softcapping = 0
|
||||
if logit_scaling is None:
|
||||
logit_scaling = getattr(config, "logit_scale", 0)
|
||||
if logit_scaling is None:
|
||||
logit_scaling = 0
|
||||
if logit_scaling == 0:
|
||||
logit_scaling = getattr(config, "logit_scaling", 0)
|
||||
if logit_scaling is None:
|
||||
logit_scaling = 0
|
||||
model_type = getattr(config, "model_type", None)
|
||||
if model_type == "granite":
|
||||
logits_scaling = getattr(config, "logits_scaling", 1)
|
||||
if logits_scaling is None:
|
||||
logits_scaling = 1
|
||||
logit_scaling = 1 / logits_scaling
|
||||
elif model_type == "falcon_h1":
|
||||
logit_scaling = getattr(config, "lm_head_multiplier", 0)
|
||||
if logit_scaling is None:
|
||||
logit_scaling = 0
|
||||
|
||||
if logit_softcapping is None:
|
||||
logit_softcapping = 0
|
||||
if logit_scaling is None:
|
||||
logit_scaling = 0
|
||||
|
||||
return logit_softcapping, logit_scaling
|
||||
|
||||
|
||||
def effective_logits(
|
||||
logits: torch.Tensor,
|
||||
model: Optional[nn.Module] = None,
|
||||
|
|
@ -118,22 +157,11 @@ def effective_logits(
|
|||
Returns:
|
||||
Transformed logits with scaling and softcapping applied.
|
||||
"""
|
||||
# Read from model config if not provided
|
||||
if model is not None:
|
||||
config = getattr(model, "config", None)
|
||||
if config is not None:
|
||||
if logit_softcapping is None:
|
||||
logit_softcapping = getattr(config, "final_logit_softcapping", 0)
|
||||
if logit_scaling is None:
|
||||
logit_scaling = getattr(config, "logit_scale", 0)
|
||||
if logit_scaling == 0:
|
||||
logit_scaling = getattr(config, "logit_scaling", 0)
|
||||
|
||||
# Default to no transformation
|
||||
if logit_softcapping is None:
|
||||
logit_softcapping = 0
|
||||
if logit_scaling is None:
|
||||
logit_scaling = 0
|
||||
logit_softcapping, logit_scaling = _resolve_logit_params(
|
||||
model,
|
||||
logit_softcapping,
|
||||
logit_scaling,
|
||||
)
|
||||
|
||||
# Convert to float32 for stability
|
||||
x = logits.float()
|
||||
|
|
@ -886,14 +914,11 @@ def compute_asft_loss(
|
|||
raise ValueError(f"Unknown streaming mode: {mode}")
|
||||
|
||||
# Get model config for softcapping/scaling
|
||||
config = getattr(model, "config", None)
|
||||
logit_softcapping = 0
|
||||
logit_scaling = 0
|
||||
if config is not None:
|
||||
logit_softcapping = getattr(config, "final_logit_softcapping", 0)
|
||||
logit_scaling = getattr(config, "logit_scale", 0)
|
||||
if logit_scaling == 0:
|
||||
logit_scaling = getattr(config, "logit_scaling", 0)
|
||||
logit_softcapping, logit_scaling = _resolve_logit_params(
|
||||
model,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
|
||||
# Build forward inputs (without labels/num_items to force logits materialization)
|
||||
forward_inputs = {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue