Add kl_direction and normalize_by parameters to ASFT
Adds `kl_direction` ("forward"/"reverse") to control KL divergence computation direction and `normalize_by` ("tokens"/"weights") for DFT/ASFT loss normalization. Forward KL (default) matches original ASFT code behavior despite paper terminology. Reverse KL enables mode-seeking behavior. Updates `_compute_kl_divergence`, streaming strategies, `compute_asft_loss`, and `ASFTTrainer` to propagate both parameters. Adds tests for reverse KL computation
This commit is contained in:
Can
Daniel Han
parent
6f42444803
commit
4a941fbeb9
3 changed files with 136 additions and 20 deletions
|
|
@ -334,7 +334,7 @@ class TestKLDivergence:
|
|||
cur_logits = torch.randn(4, 8) # (B*T, V)
|
||||
ref_logits = torch.randn(4, 8)
|
||||
|
||||
kl = _compute_kl_divergence(cur_logits, ref_logits)
|
||||
kl = _compute_kl_divergence(cur_logits, ref_logits, kl_direction = "forward")
|
||||
|
||||
# KL should be non-negative
|
||||
assert torch.all(kl >= -1e-6) # Allow small numerical errors
|
||||
|
|
@ -343,7 +343,7 @@ class TestKLDivergence:
|
|||
"""Test that KL is zero when distributions are identical."""
|
||||
logits = torch.randn(4, 8)
|
||||
|
||||
kl = _compute_kl_divergence(logits, logits.clone())
|
||||
kl = _compute_kl_divergence(logits, logits.clone(), kl_direction = "forward")
|
||||
|
||||
# Should be close to zero
|
||||
assert torch.allclose(kl, torch.zeros_like(kl), atol = 1e-5)
|
||||
|
|
@ -353,11 +353,27 @@ class TestKLDivergence:
|
|||
cur_logits = torch.randn(2, 4, 8) # (B, T, V)
|
||||
ref_logits = torch.randn(2, 4, 8)
|
||||
|
||||
kl = _compute_kl_divergence(cur_logits, ref_logits)
|
||||
kl = _compute_kl_divergence(cur_logits, ref_logits, kl_direction = "forward")
|
||||
|
||||
# Should be flattened to (B*T,)
|
||||
assert kl.shape == (8,)
|
||||
|
||||
def test_kl_reverse_matches_manual(self):
|
||||
"""Test reverse KL matches manual computation."""
|
||||
torch.manual_seed(321)
|
||||
cur_logits = torch.randn(2, 5)
|
||||
ref_logits = torch.randn(2, 5)
|
||||
|
||||
kl_reverse = _compute_kl_divergence(
|
||||
cur_logits, ref_logits, kl_direction = "reverse"
|
||||
)
|
||||
|
||||
cur_p = F.softmax(cur_logits, dim = -1)
|
||||
ref_p = F.softmax(ref_logits, dim = -1)
|
||||
manual = (cur_p * (cur_p.log() - ref_p.log())).sum(dim = -1)
|
||||
|
||||
assert torch.allclose(kl_reverse, manual, atol = 1e-5)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Test DFT weights computation
|
||||
|
|
@ -439,6 +455,37 @@ class TestComputeASFTLoss:
|
|||
assert loss.dim() == 0
|
||||
assert loss.requires_grad
|
||||
|
||||
def test_dft_normalize_by_weights(self, simple_model):
|
||||
"""Test DFT normalization by weight sum."""
|
||||
inputs = {
|
||||
"input_ids": torch.tensor([[1, 2, 3, 4]]),
|
||||
"labels": torch.tensor([[1, 2, 3, 4]]),
|
||||
}
|
||||
|
||||
logits = simple_model(input_ids = inputs["input_ids"]).logits
|
||||
shift_labels = build_shift_labels(inputs["labels"])
|
||||
valid_mask = shift_labels != -100
|
||||
ce_losses, _ = fast_cross_entropy_loss_per_token(logits, shift_labels)
|
||||
ce_losses = ce_losses.view(shift_labels.shape)
|
||||
dft_weights = _compute_dft_weights(
|
||||
logits,
|
||||
shift_labels,
|
||||
ce_losses = ce_losses,
|
||||
valid_mask = valid_mask,
|
||||
).view(shift_labels.shape)
|
||||
token_loss = ce_losses * dft_weights
|
||||
expected = token_loss[valid_mask].sum() / dft_weights[valid_mask].sum().clamp_min(1e-8)
|
||||
|
||||
loss = compute_asft_loss(
|
||||
simple_model,
|
||||
inputs,
|
||||
asft_mode = "dft",
|
||||
kl_weight = 0.0,
|
||||
normalize_by = "weights",
|
||||
)
|
||||
|
||||
assert torch.allclose(loss, expected, atol = 1e-5)
|
||||
|
||||
def test_sft_kl_mode(self, simple_model):
|
||||
"""Test SFT+KL mode."""
|
||||
inputs = {
|
||||
|
|
@ -592,6 +639,7 @@ class TestStreamingModeMapping:
|
|||
logit_softcapping = 0,
|
||||
logit_scaling = 0,
|
||||
force_fp32 = True,
|
||||
kl_direction = "forward",
|
||||
):
|
||||
batch, seq_len = shift_labels.shape
|
||||
return torch.zeros(batch, seq_len, device = shift_labels.device)
|
||||
|
|
@ -726,6 +774,7 @@ class TestStreamingModeMapping:
|
|||
logit_softcapping = 0,
|
||||
logit_scaling = 0,
|
||||
force_fp32 = True,
|
||||
kl_direction = "forward",
|
||||
):
|
||||
batch, seq_len = ref_logits.shape[:2]
|
||||
return torch.zeros(batch * seq_len, device = ref_logits.device)
|
||||
|
|
@ -1144,8 +1193,10 @@ class TestASFTTrainerComputeLoss:
|
|||
trainer.asft_enabled = True
|
||||
trainer.asft_mode = "sft"
|
||||
trainer.kl_weight = 0.0
|
||||
trainer.kl_direction = "forward"
|
||||
trainer.reference_policy = "disable_adapter"
|
||||
trainer.asft_streaming = ASFTStreamingConfig()
|
||||
trainer.normalize_by = "tokens"
|
||||
trainer._asft_original_model = None
|
||||
|
||||
model = nn.Module()
|
||||
|
|
@ -1168,8 +1219,10 @@ class TestASFTTrainerComputeLoss:
|
|||
assert loss_mock.call_args.kwargs["model"] is model
|
||||
assert loss_mock.call_args.kwargs["asft_mode"] == "sft"
|
||||
assert loss_mock.call_args.kwargs["kl_weight"] == 0.0
|
||||
assert loss_mock.call_args.kwargs["kl_direction"] == "forward"
|
||||
assert loss_mock.call_args.kwargs["reference_policy"] == "disable_adapter"
|
||||
assert loss_mock.call_args.kwargs["streaming_config"] is trainer.asft_streaming
|
||||
assert loss_mock.call_args.kwargs["normalize_by"] == "tokens"
|
||||
|
||||
def test_compute_loss_creates_frozen_copy_once(self):
|
||||
"""Test frozen copy is created once when needed."""
|
||||
|
|
@ -1179,8 +1232,10 @@ class TestASFTTrainerComputeLoss:
|
|||
trainer.asft_enabled = True
|
||||
trainer.asft_mode = "asft"
|
||||
trainer.kl_weight = 0.1
|
||||
trainer.kl_direction = "forward"
|
||||
trainer.reference_policy = "frozen_copy"
|
||||
trainer.asft_streaming = ASFTStreamingConfig()
|
||||
trainer.normalize_by = "tokens"
|
||||
trainer._asft_original_model = None
|
||||
|
||||
model = nn.Module()
|
||||
|
|
@ -1190,7 +1245,7 @@ class TestASFTTrainerComputeLoss:
|
|||
"labels": torch.tensor([[1, 2, 3, 4]]),
|
||||
}
|
||||
|
||||
with patch(
|
||||
with pytest.warns(UserWarning), patch(
|
||||
"unsloth.trainer.deepcopy", return_value = model_copy
|
||||
) as deepcopy_mock, patch(
|
||||
"unsloth.trainer.compute_asft_loss",
|
||||
|
|
@ -1212,8 +1267,10 @@ class TestASFTTrainerComputeLoss:
|
|||
trainer.asft_enabled = True
|
||||
trainer.asft_mode = "asft"
|
||||
trainer.kl_weight = 0.1
|
||||
trainer.kl_direction = "forward"
|
||||
trainer.reference_policy = "disable_adapter"
|
||||
trainer.asft_streaming = ASFTStreamingConfig()
|
||||
trainer.normalize_by = "tokens"
|
||||
trainer._asft_original_model = None
|
||||
|
||||
model = MagicMock()
|
||||
|
|
|
|||
|
|
@ -335,10 +335,18 @@ def _compute_kl_divergence(
|
|||
logit_softcapping: float = 0,
|
||||
logit_scaling: float = 0,
|
||||
force_fp32: bool = True,
|
||||
kl_direction: Literal["forward", "reverse"] = "forward",
|
||||
) -> torch.Tensor:
|
||||
"""Compute per-token KL divergence: KL(p_ref || p_cur).
|
||||
"""Compute per-token KL divergence (forward KL by default).
|
||||
|
||||
KL(p_ref || p_cur) = sum_i p_ref(i) * (log p_ref(i) - log p_cur(i))
|
||||
KL naming is frequently confused due to PyTorch's kl_div signature
|
||||
(target is the weighting distribution). Definitions here follow standard
|
||||
math/RLHF convention:
|
||||
- Forward KL: KL(p_ref || p_cur), expectation over p_ref (mass-covering).
|
||||
- Reverse KL: KL(p_cur || p_ref), expectation over p_cur (mode-seeking).
|
||||
The original ASFT repo's paper text says "reverse KL" but its code uses
|
||||
F.kl_div(log(cur), ref), which is forward KL; we match the code behavior.
|
||||
Reverse KL is available via kl_direction="reverse".
|
||||
|
||||
Args:
|
||||
cur_logits: Current model logits (B*T, V) or (B, T, V).
|
||||
|
|
@ -347,6 +355,7 @@ def _compute_kl_divergence(
|
|||
logit_softcapping: Softcapping value.
|
||||
logit_scaling: Scaling value.
|
||||
force_fp32: Whether to compute in FP32 for stability.
|
||||
kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref).
|
||||
|
||||
Returns:
|
||||
Per-token KL divergence of shape (B*T,) or (B, T).
|
||||
|
|
@ -367,14 +376,19 @@ def _compute_kl_divergence(
|
|||
cur_eff = cur_eff.float()
|
||||
ref_eff = ref_eff.float()
|
||||
|
||||
# Compute log probabilities and probabilities
|
||||
cur_logp = F.log_softmax(cur_eff, dim = -1)
|
||||
ref_p = F.softmax(ref_eff, dim = -1)
|
||||
|
||||
# KL(p_ref || p_cur) = sum_i p_ref(i) * (log p_ref(i) - log p_cur(i))
|
||||
# Using F.kl_div: kl_div(input=log_cur, target=ref) computes the right thing
|
||||
# with reduction='none', we get per-element, then sum over vocab
|
||||
kl = F.kl_div(cur_logp, ref_p, reduction = "none").sum(dim = -1)
|
||||
if kl_direction == "forward":
|
||||
# Forward KL: KL(p_ref || p_cur)
|
||||
cur_logp = F.log_softmax(cur_eff, dim = -1)
|
||||
ref_p = F.softmax(ref_eff, dim = -1)
|
||||
# Using F.kl_div: kl_div(input=log_cur, target=ref) computes KL(ref || cur)
|
||||
kl = F.kl_div(cur_logp, ref_p, reduction = "none").sum(dim = -1)
|
||||
elif kl_direction == "reverse":
|
||||
# Reverse KL: KL(p_cur || p_ref)
|
||||
ref_logp = F.log_softmax(ref_eff, dim = -1)
|
||||
cur_p = F.softmax(cur_eff, dim = -1)
|
||||
kl = F.kl_div(ref_logp, cur_p, reduction = "none").sum(dim = -1)
|
||||
else:
|
||||
raise ValueError(f"Unknown kl_direction: {kl_direction}")
|
||||
|
||||
return kl
|
||||
|
||||
|
|
@ -487,6 +501,7 @@ def _compute_kl_batch_micro(
|
|||
logit_softcapping: float = 0,
|
||||
logit_scaling: float = 0,
|
||||
force_fp32: bool = True,
|
||||
kl_direction: Literal["forward", "reverse"] = "forward",
|
||||
) -> torch.Tensor:
|
||||
"""Compute KL using batch microbatching strategy.
|
||||
|
||||
|
|
@ -503,6 +518,7 @@ def _compute_kl_batch_micro(
|
|||
logit_softcapping: Softcapping value.
|
||||
logit_scaling: Scaling value.
|
||||
force_fp32: Whether to use FP32 for KL.
|
||||
kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref).
|
||||
|
||||
Returns:
|
||||
KL tensor of shape (B, T).
|
||||
|
|
@ -530,6 +546,7 @@ def _compute_kl_batch_micro(
|
|||
logit_softcapping,
|
||||
logit_scaling,
|
||||
force_fp32,
|
||||
kl_direction,
|
||||
)
|
||||
|
||||
# Reshape if needed
|
||||
|
|
@ -558,6 +575,7 @@ def _compute_kl_seq_kv_cache(
|
|||
logit_softcapping: float = 0,
|
||||
logit_scaling: float = 0,
|
||||
force_fp32: bool = True,
|
||||
kl_direction: Literal["forward", "reverse"] = "forward",
|
||||
) -> torch.Tensor:
|
||||
"""Compute KL using sequence chunking with KV cache strategy.
|
||||
|
||||
|
|
@ -577,6 +595,7 @@ def _compute_kl_seq_kv_cache(
|
|||
logit_softcapping: Softcapping value.
|
||||
logit_scaling: Scaling value.
|
||||
force_fp32: Whether to use FP32 for KL.
|
||||
kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref).
|
||||
|
||||
Returns:
|
||||
KL tensor of shape (B, T).
|
||||
|
|
@ -606,6 +625,7 @@ def _compute_kl_seq_kv_cache(
|
|||
logit_softcapping = logit_softcapping,
|
||||
logit_scaling = logit_scaling,
|
||||
force_fp32 = force_fp32,
|
||||
kl_direction = kl_direction,
|
||||
)
|
||||
if kl_mb.dim() == 1:
|
||||
mb_batch = b_end - b_start
|
||||
|
|
@ -677,6 +697,7 @@ def _compute_kl_seq_kv_cache(
|
|||
logit_softcapping,
|
||||
logit_scaling,
|
||||
force_fp32,
|
||||
kl_direction,
|
||||
)
|
||||
ref_outputs = ref_forward(**forward_inputs)
|
||||
ref_logits, _ = _unwrap_reference_outputs(ref_outputs)
|
||||
|
|
@ -687,6 +708,7 @@ def _compute_kl_seq_kv_cache(
|
|||
logit_softcapping,
|
||||
logit_scaling,
|
||||
force_fp32,
|
||||
kl_direction,
|
||||
)
|
||||
if kl_full.dim() == 1:
|
||||
kl_full = kl_full.view(batch_size, seq_len)
|
||||
|
|
@ -703,6 +725,7 @@ def _compute_kl_seq_kv_cache(
|
|||
logit_softcapping,
|
||||
logit_scaling,
|
||||
force_fp32,
|
||||
kl_direction,
|
||||
)
|
||||
|
||||
if kl_chunk.dim() == 1:
|
||||
|
|
@ -739,6 +762,7 @@ def _compute_kl_seq_kv_cache(
|
|||
logit_softcapping,
|
||||
logit_scaling,
|
||||
force_fp32,
|
||||
kl_direction,
|
||||
)
|
||||
ref_outputs = ref_forward(**forward_inputs)
|
||||
ref_logits, _ = _unwrap_reference_outputs(ref_outputs)
|
||||
|
|
@ -749,6 +773,7 @@ def _compute_kl_seq_kv_cache(
|
|||
logit_softcapping,
|
||||
logit_scaling,
|
||||
force_fp32,
|
||||
kl_direction,
|
||||
)
|
||||
if kl_full.dim() == 1:
|
||||
kl_full = kl_full.view(batch_size, seq_len)
|
||||
|
|
@ -768,9 +793,11 @@ def compute_asft_loss(
|
|||
*,
|
||||
asft_mode: Literal["sft", "dft", "sft+kl", "asft"] = "asft",
|
||||
kl_weight: float = 0.0,
|
||||
kl_direction: Literal["forward", "reverse"] = "forward",
|
||||
reference_policy: Literal["disable_adapter", "frozen_copy"] = "disable_adapter",
|
||||
streaming_config: Optional[ASFTStreamingConfig] = None,
|
||||
original_model: Optional[nn.Module] = None,
|
||||
normalize_by: Literal["tokens", "weights"] = "tokens",
|
||||
return_outputs: bool = False,
|
||||
) -> Union[torch.Tensor, Tuple[torch.Tensor, Any]]:
|
||||
"""Compute ASFT loss.
|
||||
|
|
@ -786,9 +813,11 @@ def compute_asft_loss(
|
|||
- "sft+kl": CE + KL divergence from reference
|
||||
- "asft": DFT + KL divergence (full ASFT)
|
||||
kl_weight: Weight for KL term (only for sft+kl and asft modes).
|
||||
kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref).
|
||||
reference_policy: How to get reference distribution.
|
||||
streaming_config: Configuration for streaming strategies.
|
||||
original_model: Optional pre-created frozen reference model.
|
||||
normalize_by: "tokens" (default, matches reference) or "weights" for DFT/ASFT.
|
||||
return_outputs: Whether to return model outputs alongside loss.
|
||||
|
||||
Returns:
|
||||
|
|
@ -841,10 +870,10 @@ def compute_asft_loss(
|
|||
|
||||
# Valid mask and normalization
|
||||
valid_mask = shift_labels != -100
|
||||
n_items = inputs.get("num_items_in_batch", None)
|
||||
if n_items is None:
|
||||
n_items = valid_mask.sum()
|
||||
n_items = max(n_items, 1) # Avoid division by zero
|
||||
n_items_tokens = inputs.get("num_items_in_batch", None)
|
||||
if n_items_tokens is None:
|
||||
n_items_tokens = valid_mask.sum()
|
||||
n_items_tokens = max(n_items_tokens, 1) # Avoid division by zero
|
||||
|
||||
# Handle edge case: no valid tokens
|
||||
if valid_mask.sum() == 0:
|
||||
|
|
@ -862,6 +891,7 @@ def compute_asft_loss(
|
|||
ce_losses = ce_losses.view(batch_size, seq_len)
|
||||
|
||||
# Initialize token losses
|
||||
dft_weights = None
|
||||
if asft_mode == "sft":
|
||||
# Standard SFT: just CE
|
||||
token_loss = ce_losses
|
||||
|
|
@ -911,6 +941,7 @@ def compute_asft_loss(
|
|||
logit_softcapping,
|
||||
logit_scaling,
|
||||
streaming_config.force_fp32_kl,
|
||||
kl_direction,
|
||||
)
|
||||
elif streaming_enabled and ref_strategy == "seq_kv_cache":
|
||||
seq_chunk_size = streaming_config.seq_chunk_size
|
||||
|
|
@ -935,6 +966,7 @@ def compute_asft_loss(
|
|||
logit_softcapping = logit_softcapping,
|
||||
logit_scaling = logit_scaling,
|
||||
force_fp32 = streaming_config.force_fp32_kl,
|
||||
kl_direction = kl_direction,
|
||||
)
|
||||
else:
|
||||
# Full reference forward
|
||||
|
|
@ -947,6 +979,7 @@ def compute_asft_loss(
|
|||
logit_softcapping,
|
||||
logit_scaling,
|
||||
streaming_config.force_fp32_kl,
|
||||
kl_direction,
|
||||
)
|
||||
kl = kl.view(batch_size, seq_len)
|
||||
del ref_logits
|
||||
|
|
@ -972,8 +1005,14 @@ def compute_asft_loss(
|
|||
else:
|
||||
raise ValueError(f"Unknown asft_mode: {asft_mode}")
|
||||
|
||||
# Final reduction: sum over valid tokens, divide by n_items
|
||||
loss = token_loss[valid_mask].sum() / n_items
|
||||
# Final reduction: sum over valid tokens, divide by chosen normalizer.
|
||||
normalizer = n_items_tokens
|
||||
if normalize_by == "weights" and dft_weights is not None:
|
||||
weight_sum = dft_weights[valid_mask].sum()
|
||||
normalizer = weight_sum.clamp_min(1e-8)
|
||||
elif normalize_by != "tokens":
|
||||
raise ValueError(f"Unknown normalize_by: {normalize_by}")
|
||||
loss = token_loss[valid_mask].sum() / normalizer
|
||||
|
||||
if return_outputs:
|
||||
return loss, outputs
|
||||
|
|
|
|||
|
|
@ -230,8 +230,10 @@ class ASFTTrainer(UnslothTrainer):
|
|||
asft_enabled: bool = False,
|
||||
asft_mode: Literal["sft", "dft", "sft+kl", "asft"] = "asft",
|
||||
kl_weight: float = 0.0,
|
||||
kl_direction: Literal["forward", "reverse"] = "forward",
|
||||
reference_policy: Literal["disable_adapter", "frozen_copy"] = "disable_adapter",
|
||||
asft_streaming: Optional[ASFTStreamingConfig] = None,
|
||||
normalize_by: Literal["tokens", "weights"] = "tokens",
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize ASFTTrainer.
|
||||
|
|
@ -246,10 +248,12 @@ class ASFTTrainer(UnslothTrainer):
|
|||
- "sft+kl": CE + KL divergence from reference
|
||||
- "asft": Full ASFT (DFT + KL)
|
||||
kl_weight: Weight for KL term (used in sft+kl and asft modes).
|
||||
kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref).
|
||||
reference_policy: How to compute reference distribution:
|
||||
- "disable_adapter": Use model with LoRA adapters disabled
|
||||
- "frozen_copy": Use a frozen deepcopy of the model
|
||||
asft_streaming: Optional streaming config for VRAM reduction.
|
||||
normalize_by: "tokens" (default) or "weights" for DFT/ASFT normalization.
|
||||
**kwargs: Keyword arguments for parent trainer.
|
||||
"""
|
||||
super().__init__(*args, **kwargs)
|
||||
|
|
@ -257,8 +261,10 @@ class ASFTTrainer(UnslothTrainer):
|
|||
self.asft_enabled = asft_enabled
|
||||
self.asft_mode = asft_mode
|
||||
self.kl_weight = kl_weight
|
||||
self.kl_direction = kl_direction
|
||||
self.reference_policy = reference_policy
|
||||
self.asft_streaming = asft_streaming or ASFTStreamingConfig()
|
||||
self.normalize_by = normalize_by
|
||||
|
||||
# Will be lazily initialized if needed
|
||||
self._asft_original_model = None
|
||||
|
|
@ -294,6 +300,18 @@ class ASFTTrainer(UnslothTrainer):
|
|||
and not hasattr(model, "disable_adapter")
|
||||
)
|
||||
if needs_frozen_copy and self._asft_original_model is None:
|
||||
if self.reference_policy == "frozen_copy":
|
||||
warnings.warn(
|
||||
"Unsloth: Creating a frozen copy of the model for ASFT. "
|
||||
"This doubles VRAM usage. Use 'disable_adapter' if using LoRA.",
|
||||
stacklevel = 2,
|
||||
)
|
||||
elif self.reference_policy == "disable_adapter":
|
||||
warnings.warn(
|
||||
"Unsloth: 'disable_adapter' is unavailable; falling back to a "
|
||||
"frozen copy for ASFT. This doubles VRAM usage.",
|
||||
stacklevel = 2,
|
||||
)
|
||||
self._asft_original_model = deepcopy(model)
|
||||
self._asft_original_model.eval()
|
||||
self._asft_original_model.requires_grad_(False)
|
||||
|
|
@ -304,9 +322,11 @@ class ASFTTrainer(UnslothTrainer):
|
|||
inputs = inputs,
|
||||
asft_mode = self.asft_mode,
|
||||
kl_weight = self.kl_weight,
|
||||
kl_direction = self.kl_direction,
|
||||
reference_policy = self.reference_policy,
|
||||
streaming_config = self.asft_streaming,
|
||||
original_model = self._asft_original_model,
|
||||
normalize_by = self.normalize_by,
|
||||
return_outputs = return_outputs,
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue