Skip to content

Add P5-evaluated model runtime and loading support - #4

Merged
rekrats merged 7 commits into
mainfrom
guozl-eval
Aug 7, 2026
Merged

Add P5-evaluated model runtime and loading support#4
rekrats merged 7 commits into
mainfrom
guozl-eval

Conversation

@rekrats

@rekrats rekrats commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR promotes the minimal reusable runtime and evaluation support exercised
during P5. It contains six issue-scoped commits plus one follow-up fix, all
rebased onto current main.

  • add Gemma 4 text/unified runtime and loader support;
  • use model-native generation boundaries and response parsing;
  • keep Qwen3.5 generation compatible with Transformers 5.4 and 5.10;
  • support the evaluated Llama-3.1-70B four-GPU runtime path;
  • thread an explicit exact base model through the evaluation runners;
  • fail closed on inconsistent Metis delta manifests/tensors;
  • rebuild populated foreign or structurally non-hybrid Qwen3.5 caches while
    preserving compatible populated hybrid caches.

Provenance of changes outside eval/

The architecture/runtime work under metis/ was not created from scratch in
this PR. The Gemma/Qwen/Llama implementation was ported as reviewed minimal diffs from code snapshots developed and exercised with the 29/30 training and backbone-transfer work, then reconciled against the current upstream main on xxx。

29/30-sourced :

  • metis/backbone_wrappers/Gemma4_wrapper.py: Gemma 4 text/unified wrapper, attention/layer handling, generation path, and hidden-state behavior;
  • metis/configuration_metis.py, metis/modeling_metis.py,
    metis/dev_beta/metis_block.py, metis/memory_utils.py, and
    metis/utils.py: Gemma-family registration, memory-layer selection, model-output/logit projection, and shared wrapper contracts;
  • metis/weight_utils.py: Gemma-family identification, Gemma text-model and
    output-head discovery, backbone weight transfer, and sliding-layer memory
    configuration;
  • metis/generation_utils.py: model-native EOS/PAD and response-boundary helpers used by the evaluated Gemma/Qwen/Llama paths;
  • metis/backbone_wrappers/Qwen3_5_wrapper.py: Transformers-5 causal-mask and hybrid-cache compatibility, while retaining the existing ZeRO-3-safe checkpoint behavior;
  • metis/backbone_wrappers/Llama_wrapper.py: Transformers-5 mask compatibility for the Llama runtime path.

File-layout rationale

  • experiment matrix/runner/policy changes stay under eval/;
  • model-family loading, wrapper, generation, and checkpoint correctness stay
    under metis/ because they are runtime capabilities used by more than one evaluation entry point;
  • no new repository test files are included, per the requested publication boundary.

@rekrats

rekrats commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

Some changes are in eval/ and have been tested.

Some changes are in metis/, including support for geema4, from machines No.29/30.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR promotes the P5-evaluated runtime/loading path into reusable metis/ and threads the necessary evaluation-runner plumbing in eval/, adding Gemma 4 (text + unified) support, tightening checkpoint correctness checks, and extending multi-GPU policy handling (notably Llama 3.1 70B).

Changes:

  • Add Gemma 4 runtime wrapper + loading/weight-copy support, including backbone-specific logit projection hooks.
  • Centralize generation boundary/decoding helpers (EOS/PAD resolution + response parsing) and apply them across evaluation scripts.
  • Harden checkpoint loading (delta manifest/tensor consistency, required tensor presence) and update model-parallel device policy (incl. 4-GPU Llama 70B), plus cache/mask compatibility for Transformers 5.x.

Reviewed changes

Copilot reviewed 20 out of 20 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
metis/weight_utils.py Add gemma4 backbone type, Gemma 4 shell loading, and Gemma-specific text model/output embedding extraction during weight copy.
metis/utils.py Introduce project_logits() hook on the backbone wrapper to support backbone-specific logit post-processing.
metis/modeling_metis.py Route logits through project_logits() and avoid returning hidden states unless explicitly requested (even when commit requires them internally).
metis/memory_utils.py Skip commit hooks for blocks without writable memory and fail fast if a captured layer lacks memory modules.
metis/generation_utils.py New shared helpers for EOS/PAD resolution and response decoding/parsing across backbones/tokenizers.
metis/dev_beta/metis_block.py Add per-layer attention-type resolution and per-layer attention geometry adaptation; allow disabling memory on sliding layers.
metis/configuration_metis.py Improve Gemma 4 config error reporting; resolve BOS/EOS/PAD IDs from outer config when needed.
metis/checkpoint_utils.py Thread new memory config fields from delta config and enforce manifest↔tensor shape/dtype/key consistency during delta load.
metis/backbone_wrappers/Qwen3_5_wrapper.py Make causal-mask creation compatible across Transformers 5.4/5.10 and harden hybrid-cache creation/rebuild logic.
metis/backbone_wrappers/Llama_wrapper.py Make causal-mask creation compatible across Transformers 5.4/5.10 by omitting removed keyword.
metis/backbone_wrappers/Gemma4_wrapper.py New Gemma 4 / Gemma 4 Unified wrapper integrating Metis blocks, sliding/full mask handling, and Gemma-specific attention + logit softcapping.
eval/methods/shared/metis_loader.py Add Gemma 4 base-model loading path, extend model-family inference + device policy (incl. Llama 70B), and improve model-parallel commit behavior and checkpoint validation.
eval/experiments/ood/run.py Thread base-model override to Metis runs and keep model-parallel options scoped to Metis.
eval/experiments/main_tables/run.py Thread explicit base-model path into the Metis memqa runner invocation.
eval/environments/paper-eval-minimal-cu118.yml Update pinned Transformers version to match Gemma 4 Unified requirements.
eval/benchmarks/ood/scripts/run_memory_only.py Add --base-model-path and pass it through to runtime loading for delta checkpoints.
eval/benchmarks/memqa/scripts/run_metis_memqa.py Use shared generation helpers and allow explicit base-model override via --model-path.
eval/benchmarks/memqa/scripts/run_base_context.py Use shared generation helpers for consistent EOS/PAD handling and response decoding.
eval/benchmarks/memops/scripts/run_qwen_plain_context.py Use shared generation helpers for consistent EOS/PAD handling and response decoding.
eval/benchmarks/memops/scripts/run_memop_memory_baseline.py Use shared generation helpers and thread base-model override through the Metis baseline loader/reporting.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread eval/methods/shared/metis_loader.py
Comment thread eval/methods/shared/metis_loader.py
Comment thread eval/methods/shared/metis_loader.py
Comment thread eval/methods/shared/metis_loader.py
@rekrats
rekrats merged commit 22f7aab into main Aug 7, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants