Add support for inference using LoRA checkpoint

Signed-off-by: Bruno Alvisio <balvisio@nvidia.com>

Author: balvisio

bionemo-recipes/recipes/evo2_megatron/README.md

@@ -395,6 +395,40 @@ rather than silently producing asymmetric behaviour.
   weights are always treated as a unit, and any asymmetric configuration will
   raise an error.
 
+### Running inference on a LoRA checkpoint
+
+A LoRA training checkpoint contains only adapter tensors — the base model weights
+are not duplicated. Point `--ckpt-dir` at the LoRA `iter_*` directory as usual:
+
+```bash
+torchrun --nproc_per_node 1 --no-python \
+  infer_evo2 \
+  --ckpt-dir /path/to/lora_run/checkpoints/iter_0000250 \
+  --prompt "ATCGATCGATCGATCG" \
+  --max-new-tokens 200
+```
+
+```bash
+torchrun --nproc_per_node 1 --no-python \
+  predict_evo2 \
+  --fasta sequences.fa \
+  --ckpt-dir /path/to/lora_run/checkpoints/iter_0000250 \
+  --output-dir ./predictions
+```
+
+When `infer_evo2` / `predict_evo2` detect a `peft` section in the checkpoint's
+`run_config.yaml`, they:
+
+1. load dense base weights from `checkpoint.pretrained_checkpoint` (the same
+   value that was supplied during LoRA training),
+2. apply the stored PEFT config (`run_config["peft"]`) to graft `LoRALinear`
+   wrappers onto the base modules,
+3. load only the adapter tensors from `--ckpt-dir`.
+
+No merge step is required. The base checkpoint referenced by
+`pretrained_checkpoint` must still exist on disk at the path recorded in
+`run_config.yaml`.
+
 ## Exporting to Vortex format
 
 Vortex is ARC Institute's inference format for Evo2 Hyena models, used by the

bionemo-recipes/recipes/evo2_megatron/src/bionemo/evo2/run/infer.py

@@ -70,7 +70,11 @@
 
 import torch
 import torch.distributed as dist
-from megatron.bridge.training.checkpointing import _load_model_weights_from_checkpoint
+from megatron.bridge.training.checkpointing import (
+    _generate_model_state_dict,
+    _load_model_weights_from_checkpoint,
+    apply_peft_adapter_filter_to_state_dict,
+)
 from megatron.bridge.training.config import DistributedInitConfig, RNGConfig
 from megatron.bridge.training.mixed_precision import get_mixed_precision_config
 from megatron.bridge.training.tokenizers.tokenizer import _HuggingFaceTokenizer
@@ -81,7 +85,7 @@
 )
 from megatron.bridge.utils.common_utils import get_world_size_safe
 from megatron.bridge.utils.instantiate_utils import instantiate
-from megatron.core import parallel_state
+from megatron.core import dist_checkpointing, parallel_state
 from megatron.core.inference.contexts import StaticInferenceContext
 from megatron.core.inference.engines.static_engine import StaticInferenceEngine
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
@@ -462,12 +466,35 @@ def setup_inference_engine(
 
     raw_model = model_provider.provide().eval().cuda()
 
-    logger.info(f"Loading weights from: {resolved_ckpt_dir}")
-    _load_model_weights_from_checkpoint(
-        checkpoint_path=str(resolved_ckpt_dir),
-        model=[raw_model],
-        dist_ckpt_strictness="ignore_all",
-    )
+    # A LoRA finetune checkpoint only contains adapter tensors; the base weights live in
+    # run_config["checkpoint"]["pretrained_checkpoint"]. Detect via the top-level `peft:`
+    # section (same signal `peft_pre_wrap_hook` uses during training).
+    peft_node = run_config.get("peft")
+    if peft_node is not None:
+        # pretrained_checkpoint may point at a training-output parent containing iter_*; resolve.
+        resolved_pretrained_dir = resolve_checkpoint_path(Path(run_config["checkpoint"]["pretrained_checkpoint"]))
+        logger.info(f"PEFT checkpoint detected. Loading base weights from: {resolved_pretrained_dir}")
+        _load_model_weights_from_checkpoint(
+            checkpoint_path=str(resolved_pretrained_dir),
+            model=[raw_model],
+            dist_ckpt_strictness="ignore_all",
+        )
+
+        logger.info("Applying PEFT adapter structure to base model")
+        peft_cfg = instantiate(peft_node)
+        raw_model = peft_cfg(raw_model, training=False)
+
+        logger.info(f"Loading adapter weights from: {resolved_ckpt_dir}")
+        sharded_sd = apply_peft_adapter_filter_to_state_dict(_generate_model_state_dict([raw_model], {}), peft_cfg)
+        loaded = dist_checkpointing.load(sharded_sd, str(resolved_ckpt_dir), strict="ignore_all")
+        raw_model.load_state_dict(loaded["model"], strict=False)
+    else:
+        logger.info(f"Loading weights from: {resolved_ckpt_dir}")
+        _load_model_weights_from_checkpoint(
+            checkpoint_path=str(resolved_ckpt_dir),
+            model=[raw_model],
+            dist_ckpt_strictness="ignore_all",
+        )
     logger.info("Weights loaded successfully")
 
     # Wrap with Float16Module

bionemo-recipes/recipes/evo2_megatron/src/bionemo/evo2/run/predict.py

@@ -69,7 +69,11 @@
 import torch
 import torch.distributed as dist
 from megatron.bridge.data.samplers import build_pretraining_data_loader
-from megatron.bridge.training.checkpointing import _load_model_weights_from_checkpoint
+from megatron.bridge.training.checkpointing import (
+    _generate_model_state_dict,
+    _load_model_weights_from_checkpoint,
+    apply_peft_adapter_filter_to_state_dict,
+)
 from megatron.bridge.training.config import DistributedInitConfig, RNGConfig
 from megatron.bridge.training.mixed_precision import MIXED_PRECISION_RECIPES, get_mixed_precision_config
 from megatron.bridge.training.tokenizers.tokenizer import _HuggingFaceTokenizer
@@ -86,7 +90,7 @@
     get_world_size_safe,
 )
 from megatron.bridge.utils.instantiate_utils import instantiate
-from megatron.core import parallel_state, tensor_parallel
+from megatron.core import dist_checkpointing, parallel_state, tensor_parallel
 from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator
 from megatron.core.tensor_parallel.mappings import _gather_along_last_dim
 from megatron.core.transformer.module import Float16Module
@@ -1117,12 +1121,36 @@ def predict(
     else:
         logger.warning("Could not determine number of layers from model structure")
 
-    logger.info(f"Loading weights from: {resolved_ckpt_dir}")
-    _load_model_weights_from_checkpoint(
-        checkpoint_path=str(resolved_ckpt_dir),
-        model=model,
-        dist_ckpt_strictness="ignore_all",
-    )
+    peft_section = run_config.get("peft")
+    if peft_section is not None:
+        pretrained_ckpt = resolve_checkpoint_path(Path(run_config["checkpoint"]["pretrained_checkpoint"]))
+        logger.info(f"Loading base model weights from: {pretrained_ckpt}")
+        _load_model_weights_from_checkpoint(
+            checkpoint_path=str(pretrained_ckpt),
+            model=model,
+            dist_ckpt_strictness="ignore_all",
+        )
+
+        unwrapped = [m.module for m in model]
+        peft_cfg = instantiate(peft_section)
+        peft_cfg(unwrapped, training=False)
+
+        logger.info(f"Loading adapter weights from: {resolved_ckpt_dir}")
+        sharded_sd = _generate_model_state_dict(unwrapped, {})
+        sharded_sd = apply_peft_adapter_filter_to_state_dict(sharded_sd, peft_cfg)
+        loaded = dist_checkpointing.load(sharded_sd, str(resolved_ckpt_dir), strict="ignore_all")
+        if len(unwrapped) == 1:
+            unwrapped[0].load_state_dict(loaded["model"], strict=False)
+        else:
+            for i, inner in enumerate(unwrapped):
+                inner.load_state_dict(loaded[f"model{i}"], strict=False)
+    else:
+        logger.info(f"Loading weights from: {resolved_ckpt_dir}")
+        _load_model_weights_from_checkpoint(
+            checkpoint_path=str(resolved_ckpt_dir),
+            model=model,
+            dist_ckpt_strictness="ignore_all",
+        )
     logger.info("Weights loaded successfully")
 
     # -------------------------------------------------------------------------