Add IS evaluator #432
Draft
Add IS evaluator #432
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Adds a new evaluator type that computes forward KL divergence by comparing student log-probs against pre-computed teacher log-probs from a HuggingFace dataset of traces. The evaluator bypasses Fast-LLM's data pipeline and loads traces directly, making it suitable for monitoring distillation quality during training. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
Replace HuggingFace wrapper with native Fast-LLM inference path: - Use InferenceRunner for forward passes instead of HF model wrapper - Create LanguageModelBatch from trace data with proper padding - Handle variable-length sequences via TokenSample lengths - Use preprocess_batch for attention mask handling This approach works for all model types including linear attention. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Add max_sequence_length config field (defaults to model's position embedding limit) - Skip traces exceeding max length with warning and count - Set global_logits=True for correct tensor-parallel behavior - Report number of skipped traces in output 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Add full support for TP, SP, PP, and DP parallelism modes - Use training's sequence_length instead of separate max_sequence_length - Use GPTBatchConfig for proper SP sequence splitting - Add HuggingFace dataset sharding for efficient DP distribution - Add all_reduce across data_group and pipeline_group - Fix device mismatch bug (move targets to GPU) - Use AttentionKwargs.sequence_first constant 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
tscholak
changed the base branch from
main
to
feature/cache-refactor-and-qwen2
Store raw logits unconditionally when global_logits=True in _logits_cross_entropy_forward_backward, fixing ForwardKL evaluation during distillation training where targets is never None. Also cleaned up ForwardKL evaluator: - Use GPTInferenceRunner instead of generic InferenceRunner - Add shuffle with configurable seed for reproducibility - Add split/seed config fields (replaced task field) - Proper padding via get_padding() and from_documents() - Remove memory tracking tooling, keep gc.collect cleanup 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Replace forward KL with importance-weighted accuracy and effective sample size - Shard by problem_id hash (not trace index) so each rank gets complete problems - Add TraceTensors dataclass with smart constructors (empty, from_traces) - Vectorize log prob computation using F.cross_entropy with completion mask - Add _scatter_logsumexp for numerically stable grouped reductions - Use allreduce_scalar for cleaner distributed reduction - Pre-tensorize all trace data for efficient batch slicing 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.