Lambent/Qwen3-4B-Base-Continued-GRPO-B
17
Experimental GRPO "continued pretraining" - rewarding the model for completions that resembled the target data (in varying complex ways). Reward is calculated differently for creative text and code.
Trained for 1034 steps on a 3090, rank 128 QLoRA with alpha 256. Learning rate 1e-6 seemed ideal.
For this one, added LLM as judge to the reward functions; gpt-4o-mini for decent reward with logprobs for continuous and gemini-3-flash-preview for difficult-to-fool binary bonus:
class RewardModel:
"""Multi-domain reward model with LLM judge for creative domains."""
def __init__(self, device: str = "cuda"):
self.device = device
print("Loading reward model components...")
self.semantic_model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
self.chrf = CHRF(word_order=2)
self.rouge = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
def compute_reward(
self,
prediction: str,
reference: str,
reward_type: str,
prefix: str = None,
) -> float:
if not prediction or not reference:
return 0.0
if reward_type == "llm_judge":
# Base reward (GT-anchored)
embs = self.semantic_model.encode([reference, prediction], convert_to_tensor=True)
sem_score = torch.nn.functional.cosine_similarity(embs[0:1], embs[1:2]).item()
chrf_score = self.chrf.sentence_score(prediction, [reference]).score / 100.0
base_reward = 0.4 * sem_score + 0.6 * chrf_score
# LLM judge bonus
if prefix is not None:
llm_reward, flash_bonus = get_llm_judge_reward(prefix, prediction, reference)
else:
llm_reward, flash_bonus = 0.5, 0.0
# Multiplicative: base * (1 + llm_bonus + flash_bonus)
# This ensures LLM bonus scales with GT similarity
return base_reward * (1 + 0.3 * llm_reward + 0.2 * flash_bonus)
elif reward_type == "creative":
# Original creative reward (fallback)
embs = self.semantic_model.encode([reference, prediction], convert_to_tensor=True)
sem_score = torch.nn.functional.cosine_similarity(embs[0:1], embs[1:2]).item()
chrf_score = self.chrf.sentence_score(prediction, [reference]).score / 100.0
return 0.4 * sem_score + 0.6 * chrf_score
elif reward_type == "hybrid":
embs = self.semantic_model.encode([reference, prediction], convert_to_tensor=True)
sem_score = torch.nn.functional.cosine_similarity(embs[0:1], embs[1:2]).item()
rouge_result = self.rouge.score(reference, prediction)
rouge_l = rouge_result['rougeL'].fmeasure
ref_len = len(reference.split())
pred_len = len(prediction.split())
if ref_len > 0:
len_ratio = pred_len / ref_len
length_penalty = max(0.0, 1.0 - abs(1.0 - len_ratio))
else:
length_penalty = 1.0 if pred_len == 0 else 0.0
return 0.6 * sem_score + 0.3 * rouge_l + 0.1 * length_penalty
elif reward_type == "levenshtein":
max_len = max(len(prediction), len(reference))
if max_len == 0:
return 1.0
dist = Levenshtein.distance(prediction, reference)
return max(0.0, 1.0 - (dist / max_len))
else:
raise ValueError(f"Unknown reward type: {reward_type}")Quick capability diagnostics compared to original (most rapid subset to test both on):