# How to Debug RWKV-CLIP Training Issues Using TensorBoard Logging and Checkpointing

> Debug RWKV-CLIP training issues effectively. Monitor dynamics with TensorBoard logging and resume experiments using epoch-wise checkpoints.

- Repository: [DeepGlint/rwkv-clip](https://github.com/deepglint/rwkv-clip)
- Tags: how-to-guide
- Published: 2026-02-28

---

**Use the built-in `SpeedCallBack` console logs and TensorBoard scalars to monitor training dynamics in real-time, while leveraging epoch-wise checkpoints saved in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) to resume experiments or inspect intermediate model states.**

The RWKV-CLIP training pipeline from deepglint/rwkv-clip provides robust instrumentation for debugging distributed training runs. By combining TensorBoard visualization with automatic model checkpointing and console speed monitoring, you can quickly diagnose convergence issues, gradient explosions, or distributed deadlocks without restarting experiments from scratch.

## Setting Up TensorBoard Logging for Real-Time Monitoring

### Initializing the SummaryWriter

The training script automatically instantiates a TensorBoard `SummaryWriter` at line 226 in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py). This writer logs scalars every `frequent` steps (default: 5), enabling real-time visualization of training dynamics.

```python

# From train.py line 226

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir=os.path.join(args.output, 'tensorboard'))

```

Scalar values are written at lines 27–33, tracking **loss**, **learning rate**, and **logit scale**—the three critical indicators for contrastive learning convergence.

### Interpreting Key Training Metrics

Launch TensorBoard from the project root to inspect the following curves:

```bash
tensorboard --logdir=output_dir/tensorboard  # replace output_dir with your --output arg

```

Monitor these specific metrics to debug training issues:

- **Loss curve** – Should exhibit smooth decay. Sudden spikes indicate learning-rate bursts or gradient explosions.
- **Learning-rate schedule** – Verify the curve matches your selected scheduler (`--lr-scheduler`). The OneCycle implementation resides at lines 61–68 in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py).
- **Logit-scale growth** – Should increase until the contrastive loss stabilizes. A constant value signals potential issues with `ClipLoss` scaling or gradient flow.

## Monitoring Training Health with SpeedCallBack

### Console Output Analysis

The `SpeedCallBack` class implemented at lines 65–92 in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) prints diagnostic information via `logging.info(msg)` at line 92. This callback reports:

- **Per-GPU speed** (`rank:{int(speed)}`) and aggregated throughput (`total:{int(speed_total)}`)
- **Estimated time of arrival** (`required:{time_for_end:.1f} hours`)
- **Current learning rate** (`lr:[{lr_1:.8f}]`)

Check these console logs to confirm that your distributed launch utilizes all ranks (verify `world_size` from environment variables) and that batch sizes propagate correctly across workers.

### Spotting Distributed Issues

If training slows dramatically after a few epochs, examine the console for **NCCL error** messages or stalled ranks. The `SpeedCallBack` helps identify when specific GPUs fall behind, often indicating memory pressure or deadlock conditions in distributed data parallel training.

## Leveraging Model Checkpoints for Debugging and Recovery

### Checkpoint Structure and Storage

At the end of each epoch, the script persists a full model state-dict using `torch.save` at line 341 in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py):

```python

# From train.py line 341

torch.save(
    obj=RWKV_CLIP_model.state_dict(),
    f=os.path.join(args.output, f"RWKV_CLIP_model_{epoch}.pt")
)

```

Each checkpoint contains the **raw state-dict** without optimizer states. The helper function `load_model_weight` in [`model/utils.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/utils.py) (line 95) handles the `module.` prefix stripping required when loading distributed training checkpoints into single-GPU inference contexts.

### Verifying Checkpoint Integrity

When loss curves behave unexpectedly, pause training and verify the latest checkpoint in isolation:

```python
import torch
from model.utils import create_RWKV_Model
from model_config.utils_notebook import load_model_configs

cfg = load_model_configs('model_config/RWKV_CLIP_B32.json')
ckpt = 'output_dir/RWKV_CLIP_model_2.pt'
model = create_RWKV_Model(cfg, model_weight_path=ckpt).cuda()
model.eval()

# Run inference sanity check

dummy_img = torch.randn(1, 3, cfg.input_size, cfg.input_size).cuda()
dummy_txt = torch.randint(0, cfg.vocab_size, (1, cfg.ctx_len)).cuda()
with torch.no_grad():
    img_feat, txt_feat, logit = model(dummy_img, dummy_txt)
print('Feature shapes:', img_feat.shape, txt_feat.shape, 'logit_scale:', logit.item())

```

To inspect specific layer weights for vanishing or exploding parameters:

```python
state = torch.load('output_dir/RWKV_CLIP_model_5.pt')
print(state['Image_RWKV.patch_embedding.weight'].abs().mean())

```

### Implementing Training Resume Functionality

While the repository does not ship with a built-in resume flag, you can add one by modifying [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py). Insert the following logic after `RWKV_CLIP_model = get_model_RWKV_CLIP(args)`:

```python

# Add to argument parser:

parser.add_argument('--resume', type=int, default=None,
                    help='Epoch number of checkpoint to resume from')

# Insert after model initialization:

if args.resume is not None:
    ckpt_path = os.path.join(args.output, f"RWKV_CLIP_model_{args.resume}.pt")
    RWKV_CLIP_model = unwrap_model(RWKV_CLIP_model)  # Access underlying nn.Module

    RWKV_CLIP_model = load_model_weight(RWKV_CLIP_model, ckpt_path)
    logging.info(f"Resumed from checkpoint {ckpt_path}")

```

## Debugging Common Training Issues

### Flat Loss and Learning Rate Anomalies

When loss remains constant across epochs, examine the TensorBoard learning-rate curve. If the schedule appears correct but loss does not decrease, check the `SpeedCallBack` output for `amp` gradient scaler values—aggressive gradient clipping or learning rates set too low/high often cause stagnation.

### NaN Loss and FP16 Instability

**NaN values** in TensorBoard scalars indicate FP16 overflow or `logit_scale` exceeding its clamp bounds. When this occurs:

1. Reduce the `GradScaler` initialization scale
2. Switch to BF16 precision using `--precision bf16`
3. Inspect the checkpoint immediately preceding the NaN event using `torch.load` to identify which layers exploded first

### Checkpoint Loading Failures

If resuming training fails with key mismatches, verify that the model architecture matches the checkpoint. Architecture changes (e.g., modifying `--image-depth`) alter the state-dict keys. Use `torch.load` to examine missing or unexpected keys, then compare against `model.state_dict().keys()` to identify discrepancies.

## Step-by-Step Debugging Workflow

Follow this systematic approach to resolve training deviations using the RWKV-CLIP instrumentation:

1. **Launch with TensorBoard enabled** – Start training with `--output my_run` to activate logging.
2. **Monitor console output** – Watch `SpeedCallBack` prints every 5 steps for speed, ETA, and loss values.
3. **Validate TensorBoard curves** – Confirm loss decreases, learning-rate follows the expected schedule, and logit-scale grows appropriately.
4. **Inspect checkpoints on anomalies** – When curves deviate, load the latest `RWKV_CLIP_model_<epoch>.pt` in a REPL to run inference sanity checks.
5. **Analyze weight statistics** – Use `torch.load` to compare layer norms across epochs and identify vanishing or exploding parameters.
6. **Resume with corrections** – Adjust hyperparameters (e.g., lower `--lr`, increase `--gradient-acc`, switch `--precision`) and restart from the stable checkpoint using your custom `--resume` flag.
7. **Iterate** – Repeat until the loss curve exhibits stable convergence.

## Summary

- **TensorBoard integration** in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) (lines 27–33, 226) provides real-time visualization of loss, learning-rate, and logit-scale metrics essential for diagnosing convergence issues.
- **SpeedCallBack** (lines 65–92) delivers console diagnostics including per-GPU throughput, ETA, and current learning rate to detect distributed training stalls.
- **Automatic checkpointing** at line 341 saves raw state-dicts every epoch, enabling training resumption and intermediate weight inspection without restarting from scratch.
- **Checkpoint verification** using `create_RWKV_Model` and `load_model_weight` (from [`model/utils.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/utils.py) line 95) allows isolated testing of saved weights to confirm model integrity.
- **Common debugging patterns** include monitoring for NaN loss (indicating FP16 overflow), flat loss curves (signaling learning-rate issues), and key mismatches when loading checkpoints across different architectures.

## Frequently Asked Questions

### How do I resume training from a specific epoch in RWKV-CLIP?

The repository does not provide a native `--resume` flag, but you can implement one by adding a `resume` argument to the parser and loading the checkpoint after model initialization. Use `unwrap_model()` to access the underlying module, then call `load_model_weight()` from [`model/utils.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/utils.py) to strip the `module.` prefix and load the state-dict before the training loop begins.

### Why does my loss show NaN values in TensorBoard after a few epochs?

NaN loss typically indicates FP16 gradient overflow or the `logit_scale` parameter exceeding its clamp bounds. Check the TensorBoard scalars immediately before the NaN appears—if `logit_scale` spikes, reduce the initial scale of the `GradScaler` or switch to BF16 precision using `--precision bf16`. You can also inspect the checkpoint from the previous epoch to identify which layers exhibited exploding gradients.

### How can I verify that a saved checkpoint is not corrupted?

Load the checkpoint in an isolated Python session using `torch.load()`, then initialize a model via `create_RWKV_Model()` with the appropriate config file. Pass the checkpoint path to `model_weight_path` and run a forward pass with dummy image and text tensors. If the forward pass executes without errors and produces expected feature shapes, the checkpoint is valid.

### What should I check when training speed suddenly drops across all GPUs?

Examine the `SpeedCallBack` console output for "NCCL error" messages or verify that each rank writes to its own `training.log` file. Sudden speed drops often indicate distributed barrier deadlocks or out-of-memory conditions on specific ranks. Confirm that the `speed` and `speed_total` values reported in the console remain consistent across steps—discrepancies suggest that one or more ranks have stalled.