fix: resolve audited training and inference bugs

- reject prompts that encode to zero tokens in add_task instead of admitting a task whose prefill can never run, and surface empty-id run_batch calls as prompt_empty errors
- deliver the STOP stream callback when cancelling a live task so clients observe termination instead of hanging until socket timeout
- strip the torch.compile _orig_mod. prefix at every unwrap_model site and when loading checkpoints so FSDP state dicts and saved weights no longer leak the wrapper name into downstream keys
- reject online_* train strategies with nprocs > 1 at config validation time, explaining the NCCL all-gather deadlock they would otherwise hit mid-run
- apply the frequency penalty before temperature scaling (OpenAI semantics) so the penalty survives temperature=0 instead of being annihilated by the 1e8 logit blowup, and exclude penalty pipelines from the greedy fast path
- return logprobs from the raw pre-strategy distribution so they match training-side policy logprobs for PPO/GRPO importance ratios

astrai/trainer/train_context.py

@@ -13,7 +13,12 @@ from astrai.config.train_config import TrainConfig
 from astrai.dataset import RDSampler
 from astrai.inference.scheduler import InferenceScheduler
 from astrai.model.components.lora import inject_lora
-from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model
+from astrai.parallel.executor import (
+    BaseExecutor,
+    ExecutorFactory,
+    create_ref_model,
+    strip_compile_prefix,
+)
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
 from astrai.protocols import OptimizerProtocol, SchedulerProtocol
 from astrai.serialization import (
@@ -145,6 +150,7 @@ class TrainContextBuilder:
                         checkpoint.state_dict,
                         ConfigFactory.load(checkpoint.config or state.model_config),
                     )
+                checkpoint.state_dict = strip_compile_prefix(checkpoint.state_dict)
                 state.state_dict = checkpoint.state_dict
                 state.model_config = checkpoint.config or state.model_config
                 if self._resume:
@@ -192,7 +198,16 @@ class TrainContextBuilder:
                     target_modules=set(cfg.lora.target_modules),
                 )
             if state.state_dict is not None:
-                model.load_state_dict(state.state_dict, strict=False)
+                result = model.load_state_dict(state.state_dict, strict=False)
+                if result.missing_keys or result.unexpected_keys:
+                    logger.warning(
+                        "preloaded state dict mismatch: %d missing, %d unexpected "
+                        "(first missing: %s, first unexpected: %s)",
+                        len(result.missing_keys),
+                        len(result.unexpected_keys),
+                        result.missing_keys[:3],
+                        result.unexpected_keys[:3],
+                    )
             return model
 
         def after_wrap(model):