How to Convert and Run RWKV-CLIP in Different Precision Modes (fp32, fp16, bf16)

RWKV-CLIP supports three numeric precision modes—fp32, fp16, and bf16—controlled via the --precision command-line argument and the RWKV_FLOAT_MODE environment variable, with automatic dtype casting implemented in model/Text_rwkv.py and mixed-precision training logic in train.py.

The deepglint/rwkv-clip repository provides flexible precision configuration to balance memory usage, training speed, and numerical stability. Whether you are running zero-shot inference on consumer GPUs or training large-scale vision-language models, selecting the appropriate precision mode is critical for optimal performance.

Understanding RWKV-CLIP Precision Modes

RWKV-CLIP implements three distinct floating-point formats that trade off between computational efficiency and numerical precision.

fp32 (Full Precision)

fp32 uses torch.float32 (32-bit floating point) and serves as the default mode for inference scripts. This mode provides the highest numerical accuracy but consumes the most GPU memory and offers the slowest throughput. In zero_shot.py, fp32 is the default when --precision is omitted.

fp16 (Half Precision)

fp16 utilizes torch.float16 (16-bit floating point) to reduce memory footprint by approximately 50% compared to fp32. When selected, the model casts weights using .half() in model/Text_rwkv.py. However, fp16 requires gradient scaling during training to prevent underflow, as implemented in train.py using torch.cuda.amp.GradScaler.

bf16 (Brain Floating Point)

bf16 employs torch.bfloat16, offering the same memory savings as fp16 but with a wider dynamic range that improves training stability. In train.py, bf16 is the default precision (--precision bf16), and it does not require gradient scaling. The conversion occurs via .bfloat16() in model/Text_rwkv.py.

How Precision Selection Works in the Code

The precision configuration flows through the codebase via command-line arguments, environment variables, and dtype casting operations.

Command-Line Argument Parsing

Inference and training scripts expose the --precision argument to users. In zero_shot.py (lines 334-363), the argument is defined as:

parser.add_argument("--precision", default="fp32", type=str)

Similarly, train.py (lines 158-166) sets the default to bf16:

parser.add_argument("--precision", default="bf16", type=str)

Environment Variable Propagation

After parsing, the selected precision is stored in the RWKV_FLOAT_MODE environment variable to ensure consistency across modules. In zero_shot.py (line 362):

os.environ['RWKV_FLOAT_MODE'] = str(args.precision)

In train.py (lines 13-16), this occurs early in the script:

os.environ['RWKV_FLOAT_MODE'] = str(args.precision)

Weight Casting in Text_rwkv.py

The model/Text_rwkv.py file (lines 71-76) reads the environment variable and casts model weights accordingly:

if os.environ["RWKV_FLOAT_MODE"] == "fp16":
    m[n] = m[n].half()
elif os.environ["RWKV_FLOAT_MODE"] == "bf16":
    m[n] = m[n].bfloat16()

This ensures that all model parameters are converted to the target dtype before training or inference begins.

Training with Different Precision Modes

The training script train.py handles mixed-precision training differently depending on whether fp16 or bf16 is selected.

Autocast Configuration

During the forward pass, train.py (lines 30-34) configures the automatic mixed-precision context based on the selected precision:

training_precision_type = torch.bfloat16 if args.precision == "bf16" else torch.float16
with torch.cuda.amp.autocast(True, dtype=training_precision_type):
    # forward pass logic

This enables automatic casting of activations to the appropriate 16-bit format while maintaining 32-bit master weights for stability.

Gradient Scaling Differences

The backward pass handling diverges based on precision type (lines 4-12). For bf16, gradient scaling is unnecessary:

if args.precision == "bf16":
    loss.backward()
    optimizer.step()
else:
    # fp16 requires gradient scaler

    auto_scaler.scale(loss).backward()
    auto_scaler.step(optimizer)
    auto_scaler.update()

For fp16 training, torch.cuda.amp.GradScaler is essential to prevent gradient underflow, while bf16's wider dynamic range eliminates this requirement.

Inference with Different Precision Modes

Inference scripts apply the same precision casting mechanism but focus on memory efficiency and throughput.

Zero-Shot Evaluation

The zero_shot.py script supports all three precision modes for zero-shot image classification. When running with fp16 or bf16, the model loads with reduced memory footprint:

python zero_shot.py --precision fp16 \
    --image-path ./data/example.jpg \
    --text "a photo of a cat"

The environment variable RWKV_FLOAT_MODE ensures that model/Text_rwkv.py casts weights before the forward pass.

Text-Image Retrieval

For text-image retrieval tasks, text_image_retrieval.py implements identical precision handling. Using bf16 provides the best balance of speed and stability for large-scale retrieval:

python text_image_retrieval.py --precision bf16 \
    --image-list ./images.txt \
    --text-list ./queries.txt

Code Examples

Command-Line Examples

Run zero-shot evaluation in different precision modes:

fp32 (default):

python zero_shot.py --precision fp32 \
    --image-path ./data/example.jpg \
    --text "a photo of a cat"

fp16:

python zero_shot.py --precision fp16 \
    --image-path ./data/example.jpg \
    --text "a photo of a cat"

bf16:

python zero_shot.py --precision bf16 \
    --image-path ./data/example.jpg \
    --text "a photo of a cat"

Training with different precisions:

bf16 (default):

python train.py \
    --train-data /path/to/train \
    --train-num-samples 1000000 \
    --precision bf16

fp16:

python train.py \
    --train-data /path/to/train \
    --train-num-samples 1000000 \
    --precision fp16

Jupyter Notebook Integration

To switch precision modes programmatically:

import os
import torch
from model import get_model_RWKV_CLIP

# Select precision

precision = "bf16"  # Options: "fp32", "fp16", "bf16"

os.environ["RWKV_FLOAT_MODE"] = precision

# Load model (weights cast automatically)

model = get_model_RWKV_CLIP(args)
model.eval()

# Verify dtype

print(f"Model loaded with precision: {precision}")

Summary

• RWKV-CLIP supports three precision modes: fp32 (full precision), fp16 (half precision), and bf16 (brain floating point), selectable via the --precision argument.
• The environment variable RWKV_FLOAT_MODE propagates the selected mode to model/Text_rwkv.py, where weights are cast using .half() for fp16 or .bfloat16() for bf16.
• Training requires different handling: bf16 uses standard loss.backward() without gradient scaling, while fp16 requires torch.cuda.amp.GradScaler to prevent underflow, as implemented in train.py.
• Inference scripts (zero_shot.py, text_image_retrieval.py) use the same precision mechanism, allowing memory-efficient deployment on GPUs with limited VRAM.

Frequently Asked Questions

What is the default precision mode for RWKV-CLIP training?

The default precision for training is bf16 (brain floating point). In train.py, the --precision argument defaults to "bf16" (lines 158-166), and the script sets os.environ['RWKV_FLOAT_MODE'] = "bf16" early in execution (lines 13-16).

Does bf16 require specific GPU hardware?

Yes, bf16 requires NVIDIA Ampere architecture or newer (A100, RTX 30 series, RTX 40 series, etc.) or AMD MI200 series GPUs. While the code will attempt to run bf16 on older hardware, it may fail or fall back to fp32. fp16 is supported on older GPUs (Pascal architecture and newer).

How do I switch precision modes in a Jupyter notebook?

Set the RWKV_FLOAT_MODE environment variable before importing the model classes:

import os
os.environ["RWKV_FLOAT_MODE"] = "bf16"  # or "fp16", "fp32"

from model import get_model_RWKV_CLIP
model = get_model_RWKV_CLIP(args)

The model/Text_rwkv.py file checks this variable during weight initialization and applies .half() or .bfloat16() accordingly (lines 71-76).

Why does fp16 use gradient scaling while bf16 does not?

fp16 has a narrower dynamic range than bf16, making it susceptible to gradient underflow (values becoming zero). The train.py script (lines 4-12) uses torch.cuda.amp.GradScaler for fp16 to scale gradients before backward pass. bf16 maintains the same dynamic range as fp32 (8-bit exponent), eliminating the need for gradient scaling while retaining memory savings.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}