liger-kernel-nightly 0.4.1.dev20241114041219__py3-none-any.whl → 0.4.1.dev20241114155849__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,7 +1,23 @@
1
+ from abc import abstractmethod
2
+ from functools import partial
3
+
1
4
  import torch
5
+ from torch.nn import functional as F
2
6
 
3
7
 
4
8
  class LigerFusedLinearPreferenceBase(torch.autograd.Function):
9
+
10
+ @abstractmethod
11
+ def preference_loss_fn(chosen_logps, rejected_logps, beta=0.1):
12
+ """
13
+ Compute preference loss.
14
+ Args:
15
+ chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
16
+ rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
17
+ beta (float): Weight for the odds ratio loss.
18
+ """
19
+ raise NotImplementedError("Preference loss function must be implemented.")
20
+
5
21
  @staticmethod
6
22
  def forward(
7
23
  ctx,
@@ -11,6 +27,9 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
11
27
  bias=None,
12
28
  loss_fn=None,
13
29
  chunk_size=1,
30
+ compute_nll_loss=True,
31
+ ignore_index=-100,
32
+ beta=0.1,
14
33
  compiled=True,
15
34
  ):
16
35
  """
@@ -24,6 +43,9 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
24
43
  bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
25
44
  loss_fn (callable): Loss function to compute the loss on a chunk of input/target.
26
45
  chunk_size (int): Size of a chunk (# of batches of stacked chosen and rejected inputs).
46
+ compute_nll_loss (bool): Whether to compute NLL loss.
47
+ ignore_index (int): Index to ignore for loss computation.
48
+ beta (float): Weight for the odds ratio loss.
27
49
  compiled (bool): Whether to use torch compile for chunk accumulation.
28
50
  """
29
51
  # TODO: Tune CHUNK_SIZE to fully utilize the GPU
@@ -36,13 +58,23 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
36
58
  loss_acc = torch.zeros((), device=_input.device)
37
59
 
38
60
  chunks = max(1, _input.shape[0] // (2 * CHUNK_SIZE))
61
+ loss_func_to_call = partial(
62
+ LigerFusedLinearPreferenceBase._compute_loss,
63
+ preference_loss_fn=loss_fn,
64
+ ignore_index=ignore_index,
65
+ beta=beta,
66
+ compute_nll_loss=compute_nll_loss,
67
+ full_target=target,
68
+ )
39
69
 
40
70
  def accumulate_chunk(input_chunk, target_chunk):
41
71
  if bias is not None:
42
72
  (chunk_grad_input, chunk_grad_weight, chunk_grad_bias), (
43
73
  chunk_loss,
44
74
  (chunk_or_loss, chunk_chosen_logps, chunk_rejected_logps),
45
- ) = torch.func.grad_and_value(loss_fn, argnums=(0, 1, 3), has_aux=True)(
75
+ ) = torch.func.grad_and_value(
76
+ loss_func_to_call, argnums=(0, 1, 3), has_aux=True
77
+ )(
46
78
  input_chunk, weight, target_chunk, bias
47
79
  )
48
80
  grad_bias.add_(chunk_grad_bias)
@@ -50,7 +82,9 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
50
82
  (chunk_grad_input, chunk_grad_weight), (
51
83
  chunk_loss,
52
84
  (chunk_or_loss, chunk_chosen_logps, chunk_rejected_logps),
53
- ) = torch.func.grad_and_value(loss_fn, argnums=(0, 1), has_aux=True)(
85
+ ) = torch.func.grad_and_value(
86
+ loss_func_to_call, argnums=(0, 1), has_aux=True
87
+ )(
54
88
  input_chunk, weight, target_chunk
55
89
  )
56
90
  grad_weight.add_(chunk_grad_weight)
@@ -105,3 +139,68 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
105
139
  grad_bias = grad_bias * grad_output if grad_bias is not None else None
106
140
 
107
141
  return grad_input, grad_weight, None, grad_bias, None, None, None
142
+
143
+ @staticmethod
144
+ def _compute_loss(
145
+ input_chunk,
146
+ weight,
147
+ target_chunk,
148
+ bias=None,
149
+ preference_loss_fn=None,
150
+ full_target=None,
151
+ ignore_index=-100,
152
+ beta=0.1,
153
+ compute_nll_loss=True,
154
+ **loss_kwargs,
155
+ ):
156
+ """
157
+ Compute the total loss for a chunk of input and target, while using an alignment/preference loss function.
158
+ Args:
159
+ preference_loss_fn (callable): Loss function to compute the loss on a chunk of input/target.
160
+ input_chunk (torch.Tensor): Chunk of input tensor. Shape: (2 * chunk_size, sequence_length, hidden_size).
161
+ weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size).
162
+ target_chunk (torch.Tensor): Chunk of target tensor. Shape: (2 * chunk_size, sequence_length).
163
+ bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
164
+ full_target (torch.Tensor): Full target tensor. Shape: (batch_size, sequence_length).
165
+ ignore_index (int): Index to ignore for loss computation.
166
+ beta (float): Weight for the odds ratio loss.
167
+ loss_kwargs (dict): Additional arguments for the loss function.
168
+ """
169
+ len_chosen_chunk = target_chunk.shape[0] // 2
170
+
171
+ logits_chunk = input_chunk @ weight.t() # chunk_size x V
172
+ if bias is not None:
173
+ logits_chunk = logits_chunk + bias
174
+ log_probs_chunk = F.log_softmax(logits_chunk.float(), dim=-1)
175
+
176
+ chosen_nll_loss = 0.0
177
+ if compute_nll_loss:
178
+ chosen_nll_loss = F.nll_loss(
179
+ log_probs_chunk[:len_chosen_chunk].view(-1, log_probs_chunk.shape[-1]),
180
+ target_chunk[:len_chosen_chunk].view(-1),
181
+ reduction="sum",
182
+ ignore_index=ignore_index,
183
+ )
184
+ chosen_nll_loss = (
185
+ chosen_nll_loss
186
+ / (full_target[: full_target.shape[0] // 2] != ignore_index).sum()
187
+ )
188
+
189
+ loss_mask = target_chunk != ignore_index
190
+ label_chunk = torch.where(loss_mask, target_chunk, 0)
191
+
192
+ per_token_logps = log_probs_chunk.gather(-1, label_chunk.unsqueeze(-1)).squeeze(
193
+ -1
194
+ )
195
+ average_log_prob = (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)
196
+
197
+ chosen_logps = average_log_prob[:len_chosen_chunk]
198
+ rejected_logps = average_log_prob[len_chosen_chunk:]
199
+
200
+ alignment_loss = preference_loss_fn(
201
+ chosen_logps, rejected_logps, beta=beta, **loss_kwargs
202
+ )
203
+ alignment_loss = alignment_loss / (full_target.shape[0] // 2)
204
+
205
+ loss = chosen_nll_loss - alignment_loss
206
+ return loss, (alignment_loss, chosen_logps, rejected_logps)
@@ -1,5 +1,3 @@
1
- from functools import partial
2
-
3
1
  import torch
4
2
  import torch.nn.functional as F
5
3
 
@@ -8,79 +6,24 @@ from liger_kernel.chunked_loss.fused_linear_preference import (
8
6
  )
9
7
 
10
8
 
11
- def odds_ratio_loss(chosen_logps, rejected_logps, beta=0.1):
12
- """
13
- Compute odds-ratio loss.
14
- Args:
15
- chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
16
- rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
17
- beta (float): Weight for the odds ratio loss.
18
- """
19
- log_odds = (chosen_logps - rejected_logps) - (
20
- torch.log1p(-torch.exp(chosen_logps)) - torch.log1p(-torch.exp(rejected_logps))
21
- )
22
- ratio = F.logsigmoid(log_odds)
23
- return beta * ratio.sum()
24
-
25
-
26
- def _compute_orpo_loss(
27
- input_chunk,
28
- weight,
29
- target_chunk,
30
- bias=None,
31
- full_target=None,
32
- ignore_index=-100,
33
- beta=0.1,
34
- compute_nll_loss=True,
35
- ):
36
- """
37
- Compute ORPO loss for a chunk of input and target.
38
- Args:
39
- input_chunk (torch.Tensor): Chunk of input tensor. Shape: (2 * chunk_size, sequence_length, hidden_size).
40
- weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size).
41
- target_chunk (torch.Tensor): Chunk of target tensor. Shape: (2 * chunk_size, sequence_length).
42
- bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
43
- full_target (torch.Tensor): Full target tensor. Shape: (batch_size, sequence_length).
44
- ignore_index (int): Index to ignore for loss computation.
45
- beta (float): Weight for the odds ratio loss.
46
- """
47
- len_chosen_chunk = target_chunk.shape[0] // 2
48
-
49
- logits_chunk = input_chunk @ weight.t() # chunk_size x V
50
- if bias is not None:
51
- logits_chunk = logits_chunk + bias
52
- log_probs_chunk = F.log_softmax(logits_chunk.float(), dim=-1)
9
+ class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
53
10
 
54
- chosen_nll_loss = 0.0
55
- if compute_nll_loss:
56
- chosen_nll_loss = F.nll_loss(
57
- log_probs_chunk[:len_chosen_chunk].view(-1, log_probs_chunk.shape[-1]),
58
- target_chunk[:len_chosen_chunk].view(-1),
59
- reduction="sum",
60
- ignore_index=ignore_index,
61
- )
62
- chosen_nll_loss = (
63
- chosen_nll_loss
64
- / (full_target[: full_target.shape[0] // 2] != ignore_index).sum()
11
+ @staticmethod
12
+ def preference_loss_fn(chosen_logps, rejected_logps, beta=0.1):
13
+ """
14
+ Compute odds-ratio loss.
15
+ Args:
16
+ chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
17
+ rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
18
+ beta (float): Weight for the odds ratio loss.
19
+ """
20
+ log_odds = (chosen_logps - rejected_logps) - (
21
+ torch.log1p(-torch.exp(chosen_logps))
22
+ - torch.log1p(-torch.exp(rejected_logps))
65
23
  )
24
+ ratio = F.logsigmoid(log_odds)
25
+ return beta * ratio.sum()
66
26
 
67
- loss_mask = target_chunk != ignore_index
68
- label_chunk = torch.where(loss_mask, target_chunk, 0)
69
-
70
- per_token_logps = log_probs_chunk.gather(-1, label_chunk.unsqueeze(-1)).squeeze(-1)
71
- average_log_prob = (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)
72
-
73
- chosen_logps = average_log_prob[:len_chosen_chunk]
74
- rejected_logps = average_log_prob[len_chosen_chunk:]
75
-
76
- or_loss = odds_ratio_loss(chosen_logps, rejected_logps, beta=beta)
77
- or_loss = or_loss / (full_target.shape[0] // 2)
78
-
79
- loss = chosen_nll_loss - or_loss
80
- return loss, (or_loss, chosen_logps, rejected_logps)
81
-
82
-
83
- class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
84
27
  @staticmethod
85
28
  def forward(
86
29
  ctx,
@@ -98,15 +41,18 @@ class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
98
41
  Handles both the forward and backward pass of the final linear layer with ORPO loss.
99
42
  Inspired from LigerFusedLinearCrossEntropyFunction (https://arxiv.org/abs/2410.10989) which fuses final linear layer and CE loss.
100
43
  """
101
- orpo_loss_fn = partial(
102
- _compute_orpo_loss,
103
- full_target=target,
44
+
45
+ return LigerFusedLinearPreferenceBase.forward(
46
+ ctx=ctx,
47
+ _input=_input,
48
+ weight=weight,
49
+ target=target,
50
+ bias=bias,
51
+ loss_fn=LigerFusedLinearORPOFunction.preference_loss_fn,
52
+ compute_nll_loss=compute_nll_loss,
104
53
  ignore_index=ignore_index,
105
54
  beta=beta,
106
- compute_nll_loss=compute_nll_loss,
107
- )
108
- return LigerFusedLinearPreferenceBase.forward(
109
- ctx, _input, weight, target, bias, loss_fn=orpo_loss_fn
55
+ compiled=compiled,
110
56
  )
111
57
 
112
58
  @staticmethod
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: liger_kernel_nightly
3
- Version: 0.4.1.dev20241114041219
3
+ Version: 0.4.1.dev20241114155849
4
4
  Summary: Efficient Triton kernels for LLM Training
5
5
  License: BSD 2-CLAUSE LICENSE
6
6
  Copyright 2024 LinkedIn Corporation
@@ -1,7 +1,7 @@
1
1
  liger_kernel/env_report.py,sha256=jye8RvUkmhqaIshdeIpoUABoAu7FPKJUib4FnAfvkpw,1132
2
2
  liger_kernel/chunked_loss/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
3
- liger_kernel/chunked_loss/fused_linear_preference.py,sha256=WYsIW3t_DEPeC2VkKACimMI0LhZwOciLluDhKnaMmWE,4416
4
- liger_kernel/chunked_loss/orpo_loss.py,sha256=HgiWjS1ueVpi63ZGVveqwXkXY7EePfSzfbW-bbZlDxE,4220
3
+ liger_kernel/chunked_loss/fused_linear_preference.py,sha256=ayx-dmAx1TW9sThHJ_wUU1MqpZeJ4-SooGh0ZgVFlOA,8420
4
+ liger_kernel/chunked_loss/orpo_loss.py,sha256=DNifPpzGV_t3dfOPlPy2XKDM6M1Qne0kCbIPztvFY9U,2179
5
5
  liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
6
6
  liger_kernel/ops/cross_entropy.py,sha256=sfUb7-jIZp0EKXjg1DYy2Wdzw_Mg-mHmGoR5bpdm4tw,15526
7
7
  liger_kernel/ops/fused_linear_cross_entropy.py,sha256=JPiQ0TgPjtQ-3F5ovC0b5ZnBk067XUmzyNuGO3KZv44,9963
@@ -46,9 +46,9 @@ liger_kernel/transformers/model/qwen2.py,sha256=EyhSSzQOskGjSnCsKMZpd1s5IAIlHd5P
46
46
  liger_kernel/transformers/model/qwen2_vl.py,sha256=j6xAhp9AG195dsZK5f8dFYVM9uKtWApZrggT5Y08jn4,7055
47
47
  liger_kernel/triton/__init__.py,sha256=yfRe0zMb47QnqjecZWG7LnanfCTzeku7SgWRAwNVmzU,101
48
48
  liger_kernel/triton/monkey_patch.py,sha256=5BcGKTtdqeYchypBIBopGIWPx1-cFALz7sOKoEsqXJ0,1584
49
- liger_kernel_nightly-0.4.1.dev20241114041219.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
50
- liger_kernel_nightly-0.4.1.dev20241114041219.dist-info/METADATA,sha256=o4-CpcoMQuNzIpWwfEGDLui8edp-sWcFlXXZFDSHmlw,21556
51
- liger_kernel_nightly-0.4.1.dev20241114041219.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
52
- liger_kernel_nightly-0.4.1.dev20241114041219.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
53
- liger_kernel_nightly-0.4.1.dev20241114041219.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
54
- liger_kernel_nightly-0.4.1.dev20241114041219.dist-info/RECORD,,
49
+ liger_kernel_nightly-0.4.1.dev20241114155849.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
50
+ liger_kernel_nightly-0.4.1.dev20241114155849.dist-info/METADATA,sha256=n6ACmfunIN20yiOn358pB5lSQZS2b97-9bxRxHJoUEg,21556
51
+ liger_kernel_nightly-0.4.1.dev20241114155849.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
52
+ liger_kernel_nightly-0.4.1.dev20241114155849.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
53
+ liger_kernel_nightly-0.4.1.dev20241114155849.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
54
+ liger_kernel_nightly-0.4.1.dev20241114155849.dist-info/RECORD,,