feat: Add kimi k2 tau experimental config support

Author: bradhilton

src/art/dev/train.py

@@ -16,6 +16,7 @@ class TrainConfig(TypedDict, total=False):
     importance_sampling_level: Literal[
         "token", "sequence", "average", "geometric_average"
     ]
+    kimi_k2_tau: float | None
     logprob_calculation_chunk_size: int
     max_negative_advantage_importance_sampling_weight: float
     num_trajectories_learning_rate_multiplier_power: float

src/art/loss.py

@@ -75,6 +75,8 @@ def loss_fn(
         prob_ratio = torch.clamp(
             prob_ratio, max=max_negative_advantage_importance_sampling_weight
         )
+    if tau := experimental_config.get("kimi_k2_tau", None):
+        advantages -= tau * logprob_diff.detach()
     if experimental_config.get("ppo", True):
         policy_loss = -torch.min(
             prob_ratio * advantages,

src/art/serverless/backend.py

@@ -158,6 +158,7 @@ async def _train_model(
                 epsilon=dev_config.get("epsilon"),
                 epsilon_high=dev_config.get("epsilon_high"),
                 importance_sampling_level=dev_config.get("importance_sampling_level"),
+                kimi_k2_tau=dev_config.get("kimi_k2_tau"),
                 learning_rate=config.learning_rate,
                 max_negative_advantage_importance_sampling_weight=dev_config.get(
                     "max_negative_advantage_importance_sampling_weight"

src/art/serverless/client.py

@@ -56,6 +56,7 @@ class ExperimentalTrainingConfig(TypedDict, total=False):
     importance_sampling_level: (
         Literal["token", "sequence", "average", "geometric_average"] | None
     )
+    kimi_k2_tau: float | None
     learning_rate: float | None
     max_negative_advantage_importance_sampling_weight: float | None
     ppo: bool | None