How to change RPU config?

[maicoldubbio]

Hi there!

Suppose you've trained a model in a hardware-aware fashion, and saved it. As such, when you load_state_dict, RPU config will be loaded as well since load_rpu_config=True by default.

Now, let's say I want to change some parameters of the RPU config such as noise_model: how can you do that? Maybe if there's a way to get the current RPU, change it, and reload it into the tile.

So far, I've been using conver_to_analog after loading the analog model, specifying a new RPU config. However I guess this is incorrect, because things like digital scales coming from HWA and other parameters will be lost.

How can I do that? What is the most correct way?

[maljoras]

Hi @edminge,
thanks for the question. The most general way is, if you have a checkpoint file, to construct a new model with a new RPUConfig and do

analog_model.load_state_dict(state_dict, load_rpu_config=False)

However, there is also a convenient function to change the RPUConfig of an existing model via:

for tile in analog_model.analog_tiles():
     tile.to(rpu_config=rpu_config)

This loops through all the tiles (aka crossbars) and applies the new RPUConfig. Note that different tiles could in principle have different RPUConfigs to specify e.g. different noise levels etc, so the RPUConfig is a tile property and not a model property.

[maljoras]

I see, likely you need the newest master branch for this. There will be a new release soon on conda, right @kaoutar55?

[HeatPhoenix]

I see it here:

aihwkit/src/aihwkit/simulator/tiles/module.py

Line 204 in 8b75f76

def to(self, *args: Any, **kwargs: Any) -> "TileModule":

Thanks for the assistance, I'll just install latest from master branch locally and see how it goes.

[HeatPhoenix]

Follow-up question (after installing 0.9.0 from git), perhaps off-topic, perhaps not:

aihwkit.exceptions.TileModuleError: RPU config mismatch: Cannot replace InferenceRPUConfig with ReRamESPreset
What are the limitations to replacing RPU configs in this way? Can I replace a SingleRPUConfig with InferenceRPUConfig, if not, why not?

(as in, I want to train with ReRAM config, to for example have corrupted devices during training, but for evaluation I want realistic device noise models to be applied).

Similarly, doing model.program_analog_weights(noise_model=ReRamWan2022NoiseModel), in this situation will get TileWithPeriphery.program_weights() got an unexpected keyword argument 'noise_model', even when in evaluation mode.

In what way could I still get to the desired result of switching RPUConfigs (train with a ReRAM preset that lets me set corrupted devices, and inference/evaluation with an InferenceRPUConfig/noise model for realistic noise models).

Many thanks, again.

[maljoras]

InferenceRPUConfig and ReRamESPreset are fundamentally different concepts. The first is for chips that are used for inference only (including using ReRAM), which means you program a software model once and use crossbars only for inference (only performing the forward pass in-memory). The latter is a training preset, that means it actually performs the full training process in-memory (forward / backward / update) assuming NVMs. That is likely not what you want. You should always use InferenceRPUConfig but change the noise_model to e.g. the ReRam noise model (see e.g. aihwkit.inference.noise.reram). You could add corrupt devices there. If you want to train with corrupt devices, you can look at the modifiers that in some cases support corrupt device training (see https://github.com/IBM/aihwkit/blob/master/src/aihwkit/simulator/parameters/inference.py#L136).

Doing training with the aihwkit while using the InferenceRPUConfig means that you perform hardware-aware training (thus preparing the model in software for deployment on the inference-only chip). It thus assumes a perfect update / backward pass.

[HeatPhoenix]

I believe using InferenceRPUConfig, with corrupted devices and then training (as in, training but taking the corrupted devices into account while training) is the desired result. I'll give that a go using the modifiers. Thanks again.

Out of curiosity, why is InferenceRPUConfig and ReamESPreset (for example, or any other RPUConfig) split like this? Wouldn't you want (to have the option of) a realistic forward pass with this noise when doing this training on the devices themselves? Does this imply the forward pass in the ReRamESPreset both during training and during evaluation is not as realistic as when using InferenceRPUConfig?

I understand that InferenceRPUConfig, for example, does not include the ability to set specialized material settings and AIMC SGD algorithms for training, but that would just make the non-Inference only configs supersets of the Inference only configs, no?

[HeatPhoenix]

InferenceRPUConfig and ReRamESPreset are fundamentally different concepts. The first is for chips that are used for inference only (including using ReRAM), which means you program a software model once and use crossbars only for inference (only performing the forward pass in-memory). The latter is a training preset, that means it actually performs the full training process in-memory (forward / backward / update) assuming NVMs. That is likely not what you want. You should always use InferenceRPUConfig but change the noise_model to e.g. the ReRam noise model (see e.g. aihwkit.inference.noise.reram). You could add corrupt devices there. If you want to train with corrupt devices, you can look at the modifiers that in some cases support corrupt device training (see https://github.com/IBM/aihwkit/blob/master/src/aihwkit/simulator/parameters/inference.py#L136).

Doing training with the aihwkit while using the InferenceRPUConfig means that you perform hardware-aware training (thus preparing the model in software for deployment on the inference-only chip). It thus assumes a perfect update / backward pass.

Apologies the double post, to properly make use of your suggestion, I need to install 0.9.0 with CUDA-support, but installing from a local package (aka pulling from git -> pip install . or something similar) will only install the CPU version, which will toss an error when trying to use this functionality: RuntimeError: Error in weight_modifier.cpp:204 Requested WeightModifierType not implemented in CPU version.

Any tips for installing this version smoothly?

Secondly, perhaps I am misunderstanding WeightModifierType, but is it not possible to apply multiple types and multiple parameters? ex. Discretization and PCM_NOISE?

Many thanks.