Checkpoint Pipeline Troubleshooting

This guide helps diagnose and resolve common issues with the checkpoint pipeline system.

Quick Diagnostics

Component Discovery

Check what components are available:

from anemoi.training.checkpoint import ComponentCatalog

print("Available sources:", ComponentCatalog.list_sources())
print("Available loaders:", ComponentCatalog.list_loaders())
print("Available modifiers:", ComponentCatalog.list_modifiers())

Checkpoint Inspection

Inspect a checkpoint without full loading:

from anemoi.training.checkpoint import get_checkpoint_metadata
from pathlib import Path

metadata = get_checkpoint_metadata(Path("model.ckpt"))
print(f"File size: {metadata.get('file_size_mb', 0):.1f} MB")
print(f"Parameters: {metadata.get('num_parameters', 'unknown')}")

Common Issues and Solutions

Configuration Issues

Unknown checkpoint source type

Error Message:

CheckpointConfigError: Unknown checkpoint source: 'my_source'

Causes:

Typo in source type name
Source module not implemented or imported
Missing dependencies for specific source types

Solutions:

Check available sources:

from anemoi.training.checkpoint import ComponentCatalog
print(ComponentCatalog.list_sources())

Use correct _target_ path in configuration:

checkpoint_pipeline:
  stages:
    - _target_: my_module.MySource
      path: /path/to/checkpoint.ckpt

Failed to instantiate stage

Error Message:

CheckpointConfigError: Failed to instantiate pipeline stage from configuration

Causes:

Invalid _target_ path
Missing required parameters
Import errors in target module

Solutions:

Verify the target path is importable
Check all required parameters are provided

Test import manually:

# Test if your stage can be imported
from my_module import MyStage
stage = MyStage(param="value")

Environment and Dependencies

No checkpoint components discovered

Warning Message:

No sources components were discovered

Causes:

Module not yet implemented
Import errors in component modules
Missing dependencies

Solutions:

Check implementation status (some modules may not be implemented yet)

Try importing manually to check for errors:

try:
    from my_module import MySource
    print("Module imported successfully")
except ImportError as e:
    print(f"Module not available: {e}")

Enable debug logging to see import issues:

import logging
logging.getLogger("anemoi.training.checkpoint").setLevel(logging.DEBUG)

Remote downloads not working

Error Message:

ImportError: aiohttp is required for remote checkpoint downloads

Solution:

Install the remote dependencies:

pip install anemoi-training[remote]

File and Path Issues

Checkpoint file does not exist

Error Message:

CheckpointNotFoundError: Checkpoint file not found at /path/to/checkpoint.ckpt

Solutions:

Verify file path:

from pathlib import Path

path = Path("/path/to/checkpoint.ckpt")
print(f"Path exists: {path.exists()}")
print(f"Is file: {path.is_file()}")
print(f"Parent exists: {path.parent.exists()}")

Check permissions:
```
ls -la /path/to/checkpoint.ckpt
```

Use absolute paths in configuration:

checkpoint_pipeline:
  stages:
    - _target_: my_module.LocalSource
      path: /absolute/path/to/checkpoint.ckpt

Loading and Compatibility Issues

Shape mismatch during loading

Error Message:

CheckpointIncompatibleError: Unexpected key(s) in state_dict

Solutions:

Use non-strict loading (when using a custom loader):

stages:
  - _target_: my_module.WeightsOnlyLoader
    strict: false

Compare checkpoint and model keys:

import torch
from anemoi.training.checkpoint import compare_state_dicts

checkpoint = torch.load("checkpoint.ckpt", map_location="cpu")
checkpoint_dict = checkpoint.get("state_dict", checkpoint)

missing, unexpected, mismatches = compare_state_dicts(
    checkpoint_dict,
    model.state_dict()
)

print(f"Missing: {missing}")
print(f"Unexpected: {unexpected}")
print(f"Shape mismatches: {mismatches}")

Cannot extract state dict

Error Message:

CheckpointValidationError: Cannot find model state in checkpoint

Causes:

Non-standard checkpoint format
Checkpoint uses different key names

Solutions:

Inspect checkpoint structure:

import torch

checkpoint = torch.load("checkpoint.ckpt", map_location="cpu")
print("Available keys:", list(checkpoint.keys())[:10])

Use format detection:

from anemoi.training.checkpoint.formats import (
    detect_checkpoint_format,
    extract_state_dict,
)

fmt = detect_checkpoint_format("checkpoint.ckpt")
print(f"Detected format: {fmt}")

Network and Remote Source Issues

Failed to download checkpoint

Error Message:

CheckpointSourceError: Failed to download from https://...

Solutions:

Check network connectivity:
```
curl -I https://example.com/model.ckpt
```

Download manually first:

wget https://example.com/model.ckpt -O local_model.ckpt

Then use a local source stage:

checkpoint_pipeline:
  stages:
    - _target_: my_module.LocalSource
      path: ./local_model.ckpt

AWS S3 authentication failed

Error Message:

CheckpointSourceError: Access denied to s3://bucket/model.ckpt

Solutions:

Check AWS credentials:

aws configure list
aws s3 ls s3://bucket/

Set environment variables:

export AWS_ACCESS_KEY_ID=your-key-id
export AWS_SECRET_ACCESS_KEY=your-secret-key
export AWS_DEFAULT_REGION=us-east-1

Memory and Performance Issues

Out of memory during loading

Error Message:

RuntimeError: CUDA out of memory

Solutions:

Load on CPU first:

import torch

checkpoint = torch.load("checkpoint.ckpt", map_location="cpu")
model.load_state_dict(checkpoint["state_dict"])
model = model.cuda()

Use weights-only loading to skip optimizer state

Clear GPU memory:

import torch

torch.cuda.empty_cache()
print(f"GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")

Advanced Debugging

Enable Debug Logging

import logging

# Enable detailed checkpoint pipeline logging
logging.getLogger("anemoi.training.checkpoint").setLevel(logging.DEBUG)

# Enable all debug logging
logging.basicConfig(level=logging.DEBUG)

Pipeline Execution Tracking

After pipeline execution, check context metadata:

result = await pipeline.execute(context)

print("Pipeline execution metadata:")
for key, value in result.metadata.items():
    if key.startswith("stage_"):
        print(f"  {key}: {value}")

Manual Stage Testing

Test individual stages in isolation:

from anemoi.training.checkpoint import CheckpointContext

context = CheckpointContext(model=my_model)

try:
    result = await my_stage.process(context)
    print("Stage succeeded")
except Exception as e:
    print(f"Stage failed: {e}")
    import traceback
    traceback.print_exc()

Getting Help

Report Issues

When reporting issues, include:

Full error message and stack trace
Configuration file (sanitized)

Environment information:

import sys
import torch
import anemoi.training

print(f"Python: {sys.version}")
print(f"PyTorch: {torch.__version__}")
print(f"Anemoi Training: {anemoi.training.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")

Component discovery output:

from anemoi.training.checkpoint import ComponentCatalog

print("Sources:", ComponentCatalog.list_sources())
print("Loaders:", ComponentCatalog.list_loaders())
print("Modifiers:", ComponentCatalog.list_modifiers())

Debug Information Collection

def collect_debug_info():
    import sys, torch, platform

    info = {
        "python_version": sys.version,
        "pytorch_version": torch.__version__,
        "platform": platform.platform(),
        "cuda_available": torch.cuda.is_available(),
    }

    from anemoi.training.checkpoint import ComponentCatalog

    info["sources"] = ComponentCatalog.list_sources()
    info["loaders"] = ComponentCatalog.list_loaders()
    info["modifiers"] = ComponentCatalog.list_modifiers()

    return info

import json
print(json.dumps(collect_debug_info(), indent=2))