style: fix ruff and black formatting issues

- Add strict=False to zip() calls (ruff B905)
- Apply black formatting fixes

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: tourzhao

examples/knowledge_distillation/kd_loss.py

@@ -19,14 +19,15 @@ def store_topk_data(samples):
 
 
 def _get_topk_data(tokens):
-    key = tuple(tokens[:20].tolist() if hasattr(tokens, 'tolist') else tokens[:20])
+    key = tuple(tokens[:20].tolist() if hasattr(tokens, "tolist") else tokens[:20])
     return _topk_data_store.get(key)
 
 
 def sampled_kl_loss(args, batch, logits, sum_of_sample_mean):
     """Forward KL on teacher-sampled tokens (KD_TOP_K=0)."""
     _, log_probs_result = get_log_probs_and_entropy(
-        logits, args=args,
+        logits,
+        args=args,
         unconcat_tokens=batch["unconcat_tokens"],
         total_lengths=batch["total_lengths"],
         response_lengths=batch["response_lengths"],
@@ -37,7 +38,7 @@ def sampled_kl_loss(args, batch, logits, sum_of_sample_mean):
     entropy = log_probs_result.get("entropy", [])
 
     kl_terms = []
-    for s_lp, t_lp in zip(student_lps, batch["teacher_log_probs"]):
+    for s_lp, t_lp in zip(student_lps, batch["teacher_log_probs"], strict=False):
         kl_terms.append(t_lp.to(s_lp) - s_lp)
 
     loss = sum_of_sample_mean(torch.cat(kl_terms))
@@ -66,11 +67,16 @@ def _extract_response_log_probs(logits, unconcat_tokens, total_lengths, response
 def topk_kl_loss(args, batch, logits, sum_of_sample_mean):
     """Forward KL on teacher's top-K tokens with temperature scaling."""
     student_full_lps = _extract_response_log_probs(
-        logits, batch["unconcat_tokens"], batch["total_lengths"], batch["response_lengths"],
+        logits,
+        batch["unconcat_tokens"],
+        batch["total_lengths"],
+        batch["response_lengths"],
     )
 
     topk_data_list = [_get_topk_data(tokens) for tokens in batch["unconcat_tokens"]]
-    valid_data = [(s_lp, data) for s_lp, data in zip(student_full_lps, topk_data_list) if data is not None]
+    valid_data = [
+        (s_lp, data) for s_lp, data in zip(student_full_lps, topk_data_list, strict=False) if data is not None
+    ]
 
     if not valid_data:
         return sampled_kl_loss(args, batch, logits, sum_of_sample_mean)
@@ -84,7 +90,7 @@ def topk_kl_loss(args, batch, logits, sum_of_sample_mean):
         s_topk = s_lp.gather(1, t_ids)
         t_renorm = torch.log_softmax(t_lps / tau, dim=-1)
         s_renorm = torch.log_softmax(s_topk / tau, dim=-1)
-        kl_terms.append((tau ** 2) * (t_renorm.exp() * (t_renorm - s_renorm)).sum(dim=-1))
+        kl_terms.append((tau**2) * (t_renorm.exp() * (t_renorm - s_renorm)).sum(dim=-1))
 
     loss = sum_of_sample_mean(torch.cat(kl_terms))
     return loss, {"kd/loss": loss.detach()}

examples/knowledge_distillation/knowledge_distillation.py

@@ -110,13 +110,19 @@ async def _generate_rollout_async(args, data_source):
     semaphore = asyncio.Semaphore(max(getattr(args, "sglang_server_concurrency", 64), 1))
 
     async with aiohttp.ClientSession() as session:
-        generated_groups = await asyncio.gather(*(
-            asyncio.gather(*(_generate_sample(args, s, sampling_params, tokenizer, session, semaphore) for s in group))
-            for group in samples
-        ))
+        generated_groups = await asyncio.gather(
+            *(
+                asyncio.gather(
+                    *(_generate_sample(args, s, sampling_params, tokenizer, session, semaphore) for s in group)
+                )
+                for group in samples
+            )
+        )
 
     first = generated_groups[0][0]
-    logger.info(f"KD rollout: prompt={first.prompt[:80]!r}, response={first.response[:80]!r}, len={first.response_length}")
+    logger.info(
+        f"KD rollout: prompt={first.prompt[:80]!r}, response={first.response[:80]!r}, len={first.response_length}"
+    )
 
     token_count = sum(s.response_length for g in generated_groups for s in g)
     return RolloutFnTrainOutput(samples=generated_groups, metrics={"kd/token_count": token_count})
@@ -132,6 +138,7 @@ def generate_rollout(args, rollout_id, data_source, evaluation=False):
     # Store top-k data for loss function access
     if KD_TOP_K > 0:
         from examples.knowledge_distillation.kd_loss import store_topk_data
+
         store_topk_data(result.samples)
 
     if KD_SAVE_PATH:

examples/knowledge_distillation/offline_kd.py

@@ -39,8 +39,12 @@ def _load_from_jsonl(load_path, rollout_id, batch_size, num_rollouts_per_prompt)
             )
 
             if "teacher_top_k_ids" in record and "teacher_top_k_logprobs" in record:
-                assert len(record["teacher_top_k_ids"]) == resp_len, f"Sample {len(samples)}: top_k_ids length mismatch"
-                assert len(record["teacher_top_k_logprobs"]) == resp_len, f"Sample {len(samples)}: top_k_logprobs length mismatch"
+                assert (
+                    len(record["teacher_top_k_ids"]) == resp_len
+                ), f"Sample {len(samples)}: top_k_ids length mismatch"
+                assert (
+                    len(record["teacher_top_k_logprobs"]) == resp_len
+                ), f"Sample {len(samples)}: top_k_logprobs length mismatch"
                 sample.train_metadata = {
                     "teacher_top_k_ids": record["teacher_top_k_ids"],
                     "teacher_top_k_logprobs": record["teacher_top_k_logprobs"],
@@ -80,10 +84,13 @@ def generate_rollout(args, rollout_id, data_source, evaluation=False):
     # Store top-k data for loss function access
     if KD_TOP_K > 0:
         from examples.knowledge_distillation.kd_loss import store_topk_data
+
         store_topk_data(samples)
 
     first = samples[0][0]
-    logger.info(f"Offline KD: prompt={first.prompt[:80]!r}, response={first.response[:80]!r}, len={first.response_length}")
+    logger.info(
+        f"Offline KD: prompt={first.prompt[:80]!r}, response={first.response[:80]!r}, len={first.response_length}"
+    )
 
     token_count = sum(s.response_length for g in samples for s in g)
     return RolloutFnTrainOutput(samples=samples, metrics={"kd/token_count": token_count})