Mapping learned weights to PCMs

[maicoldubbio]

Hi all!

Suppose you want to map HWA-trained weights to physical PCM conductances. Such model will have a number of I/O scaling coefficients that are essential for proper inference execution. AFAIK, the overall per-column MVM product y_i can be expressed as:

where alpha is the input scaling factor, and gamma_i the per-column (possibly tile-wise too) output scale.

My question is, how can you get this factor for each analog tile? So that, once mapped to actual PCMs, you get the same results. (I might be missing some steps, in case let me know!)

Thanks

[coreylammie]

Hi @maicoldubbio,

Great question! Here is a complete example of how all the scales can be extracted and MVMs can be performed using PWM and conductance values for the first tile of example 25 (https://github.com/IBM/aihwkit/blob/master/examples/25_torch_tile_lenet5_hardware_aware.py):

import torch
import aihwkit
import functools
from 25_torch_tile_lenet5_hardware_aware import (
    populate_rpu_config,
    create_analog_network,
    load_images,
)
from aihwkit.simulator.presets import StandardHWATrainingPreset
from aihwkit.nn.conversion import convert_to_digital
import matplotlib.pyplot as plt
import numpy as np


def det_l2_error(predicted_tensor, truth_tensor, axis=None):
    def det_norm(t, axis=None):
        return np.array(torch.linalg.norm(torch.tensor(t), dim=axis))

    if type(predicted_tensor) == torch.Tensor:
        predicted_tensor = predicted_tensor.cpu().detach().numpy()

    predicted_tensor = np.array(predicted_tensor)
    if type(truth_tensor) == torch.Tensor:
        truth_tensor = truth_tensor.cpu().detach().numpy()

    truth_tensor = np.array(truth_tensor)
    with np.errstate(divide="ignore", invalid="ignore"):
        if axis is not None:
            a = det_norm(predicted_tensor - truth_tensor, axis=axis)
            b = det_norm(truth_tensor, axis=axis)
        else:
            a = det_norm((predicted_tensor - truth_tensor))
            b = det_norm(truth_tensor)

        l2_err = np.divide(a, b, out=np.zeros_like(a), where=b != 0)
        l2_err = np.mean(l2_err[np.isfinite(l2_err)])

    return l2_err


DEVICE = "cuda" if aihwkit.simulator.rpu_base.cuda.is_compiled() else "cpu"

standard_rpu_config = populate_rpu_config(StandardHWATrainingPreset())
network = create_analog_network(standard_rpu_config)
state_dict = torch.load("lenet_torch_tile_model.th")
network.load_state_dict(state_dict)
network.to(DEVICE)

# Test case for the first tile.
input_tensor_shape = (100, 25)
Gmax = 180
max_PWM = 256
input_ = torch.randn(input_tensor_shape, device=DEVICE)
reference_weights = torch.randn(network[0].analog_module.in_size, network[0].analog_module.out_size)
network[0].analog_module.set_weights(reference_weights.T)
weights_ = network[0].analog_module.get_weights(apply_weight_scaling=False)[0].T

# 1. Normalize tile weights from 0-Gmax (for each column).
tile_weights = weights_ * Gmax

# 2. Clip and scale inputs.
chip_input = torch.clamp(
    input_, -network[0].analog_module.input_range, network[0].analog_module.input_range
)
chip_input = chip_input * max_PWM / network[0].analog_module.input_range

# 3. Perform the on-chip MVM.
assert chip_input.max() == max_PWM and chip_input.min() == -max_PWM
assert tile_weights.max() == Gmax and tile_weights.min() == -Gmax
chip_output = torch.matmul(input_, tile_weights.to(DEVICE))
chip_output = chip_output / (
    Gmax / network[0].analog_module.input_range
)
chip_output *= network[0].analog_module.get_scales()


# Plot results.
plt.figure(1)
x = torch.matmul(input_, reference_weights.to(DEVICE)).flatten().cpu().detach()
y = chip_output.flatten().cpu().detach()
plt.scatter(x, y)
plt.xlabel("Floating point output")
plt.ylabel("Chip output")
print("l2 error: ", det_l2_error(y, x))
plt.savefig("out.png")

I have arbitrarily set values for Gmax and max_PWM, but you should be able to configure these according to your use case. Let me know if anything is unclear!

[maicoldubbio]

Thanks @coreylammie, it couldn't be any better. I have just one question: why are you sampling the input range from analog_module rather than analog tiles? That way you could loop through the AnalogWrapper and for each analog_tiles() you get the value. But I guess it's the same :)

[maljoras]

Hi, thanks @coreylammie for the answer. Indeed, @maicoldubbio you are right, one should use the analog_tiles generators instead of directly assessing analog_module attributes for readability, but also for support of all available analog tile modules. For instance, accessing directly the analog_module property will also result in errors for some analog modules such as array tiles. Thus, as a general rule, always use the analog_tiles generators, as one would use the module generators in pytorch.