diff --git a/src/aihwkit/simulator/configs/compounds.py b/src/aihwkit/simulator/configs/compounds.py index bcf6ea9a3..8308798f8 100644 --- a/src/aihwkit/simulator/configs/compounds.py +++ b/src/aihwkit/simulator/configs/compounds.py @@ -499,11 +499,38 @@ class TransferCompound(UnitCell): define the type of update used for each transfer event. """ + def __post_init__(self) -> None: + self._assert_nonzero_dw_min() + + def _assert_nonzero_dw_min(self) -> None: + """Reject infinite-granularity (``dw_min = 0``) sub-devices. + + Transfer compounds scale the transfer learning rate and the buffer + granularity by the sub-device weight granularity (``dw_min``). A zero + ``dw_min`` (infinite-granularity mode) collapses these to zero (or + divides by zero), which silently disables learning. Infinite + granularity is therefore only supported for plain pulsed devices, not + for devices placed inside a transfer compound. + + Raises: + ConfigError: if any ``unit_cell_device`` has ``dw_min == 0``. + """ + for device in self.unit_cell_devices: + if getattr(device, "dw_min", None) == 0: + raise ConfigError( + "dw_min=0 (infinite-granularity mode) is not supported for " + "devices inside a transfer compound: the transfer learning " + "rate and buffer granularity scale with the sub-device weight " + "granularity, so dw_min=0 would disable learning." + ) + def as_bindings(self, data_type: RPUDataType) -> Any: """Return a representation of this instance as a simulator bindings object.""" if not isinstance(self.unit_cell_devices, list): raise ConfigError("unit_cell_devices should be a list of devices") + self._assert_nonzero_dw_min() + n_devices = len(self.unit_cell_devices) transfer_parameters = parameters_to_bindings(self, data_type) @@ -775,6 +802,8 @@ def as_bindings(self, data_type: RPUDataType) -> Any: if not isinstance(self.unit_cell_devices, list): raise ConfigError("unit_cell_devices should be a list of devices") + self._assert_nonzero_dw_min() + n_devices = len(self.unit_cell_devices) if n_devices != 2: raise ConfigError("Only 2 devices supported for ChoppedTransferCompound") @@ -952,6 +981,8 @@ def as_bindings(self, data_type: RPUDataType) -> Any: if not isinstance(self.unit_cell_devices, list): raise ConfigError("unit_cell_devices should be a list of devices") + self._assert_nonzero_dw_min() + n_devices = len(self.unit_cell_devices) if n_devices != 2: raise ConfigError("Only 2 devices supported for ChoppedTransferCompound") diff --git a/src/aihwkit/simulator/configs/devices.py b/src/aihwkit/simulator/configs/devices.py index 01da44ca0..0d32118ff 100644 --- a/src/aihwkit/simulator/configs/devices.py +++ b/src/aihwkit/simulator/configs/devices.py @@ -164,7 +164,24 @@ class PulsedDevice(_PrintableMixin): """Parameter governing a power-law drift.""" dw_min: float = 0.001 - """Mean of the minimal update step sizes across devices and directions.""" + r"""Mean of the minimal update step sizes across devices and directions. + + Note: + Setting ``dw_min = 0`` activates **infinite-granularity** mode. + In this mode the stochastic pulse-train update is replaced by + its exact mean-field limit: + + .. math:: + + w \leftarrow w - \mathrm{lr} \cdot (d^\top x) \cdot q(w) + + where :math:`q(w)` is the device-specific response function + (the weight-dependent factor normally applied per pulse + coincidence), x and d are the forward and backward signals, respetively. + Cycle-to-cycle noise and ``dw_min``-related device-to-device variation, + like ``dw_min_dtod``, ``dw_min_std``, are dropped, while all other variation parameters, + like bounds, gamma, slopes are preserved. + """ dw_min_dtod: float = 0.3 """Device-to-device std deviation of ``dw_min`` (in relative units to diff --git a/src/rpucuda/cuda/pulsed_weight_updater.cu b/src/rpucuda/cuda/pulsed_weight_updater.cu index d4a540129..dcc2e982e 100644 --- a/src/rpucuda/cuda/pulsed_weight_updater.cu +++ b/src/rpucuda/cuda/pulsed_weight_updater.cu @@ -303,6 +303,44 @@ void PulsedWeightUpdater::doDirectUpdate( context_->template releaseSharedBuffer(RPU_BUFFER_OUT); } +template +template +void PulsedWeightUpdater::doInfiniteGranularityUpdate( + XInputIteratorT x_in, + DInputIteratorT d_in, + PulsedRPUDeviceCuda *pulsed_device, + T *dev_weights, + const T lr, + const int m_batch, + const bool x_trans, + const bool d_trans) { + + T *fpx_buffer = context_->template getSharedBuffer(RPU_BUFFER_IN, x_size_ * m_batch); + T *fpd_buffer = context_->template getSharedBuffer(RPU_BUFFER_OUT, d_size_ * m_batch); + T *grad_buffer = context_->template getSharedBuffer(RPU_BUFFER_TEMP_0, x_size_ * d_size_); + + const T *x_out = copyIterator2Buffer(x_in, fpx_buffer, x_size_ * m_batch); + const T *d_out = copyIterator2Buffer(d_in, fpd_buffer, d_size_ * m_batch); + + // G = -lr * D^T @ X, written into grad_buffer (beta=0 overwrites). + // Uniform gemm path; for m_batch==1 this is a rank-1 update equivalent to ger. + RPU::math::gemm( + context_, d_trans, !x_trans, + d_size_, // M + x_size_, // N + m_batch, // K + -lr, d_out, d_trans ? m_batch : d_size_, x_out, x_trans ? m_batch : x_size_, (T)0.0, + grad_buffer, d_size_); + + const int total_size = x_size_ * d_size_; + pulsed_device->doInfiniteGranularityUpdate( + dev_weights, grad_buffer, context_->getRandomStates(total_size)); + + context_->template releaseSharedBuffer(RPU_BUFFER_TEMP_0); + context_->template releaseSharedBuffer(RPU_BUFFER_OUT); + context_->template releaseSharedBuffer(RPU_BUFFER_IN); +} + template bool PulsedWeightUpdater::checkForFPUpdate( AbstractRPUDeviceCuda *rpucuda_device_in, const PulsedUpdateMetaParameter &up) { @@ -358,6 +396,16 @@ void PulsedWeightUpdater::update( // safe because of isPulsedDevice PulsedRPUDeviceCudaBase *rpucuda_device = static_cast *>(rpucuda_device_in); + + // Infinite granularity mode (dw_min=0): exact mean-field update + if (rpucuda_device->getWeightGranularity() <= (T)0.0) { + auto *pulsed_device = static_cast *>(rpucuda_device); + T abs_lr = lr < (T)0.0 ? -lr : lr; + doInfiniteGranularityUpdate( + x_in, d_in, pulsed_device, dev_weights, abs_lr, m_batch, x_trans, d_trans); + return; + } + bool force_tuning = false; // check need for init (or re-init) @@ -428,6 +476,9 @@ void PulsedWeightUpdater::update( template void PulsedWeightUpdater::doDirectUpdate( \ XITERT, DITERT, AbstractRPUDeviceCuda *, NUM_T *, const NUM_T, \ const PulsedUpdateMetaParameter &, const int, const bool, const bool, const NUM_T); \ + template void PulsedWeightUpdater::doInfiniteGranularityUpdate( \ + XITERT, DITERT, PulsedRPUDeviceCuda *, NUM_T *, const NUM_T, const int, const bool, \ + const bool); \ template void PulsedWeightUpdater::tuneUpdate( \ pwukp_t &, pwukpvec_t &, XITERT, DITERT, NUM_T *, \ PulsedRPUDeviceCudaBase *, const PulsedUpdateMetaParameter &, const NUM_T, \ diff --git a/src/rpucuda/cuda/pulsed_weight_updater.h b/src/rpucuda/cuda/pulsed_weight_updater.h index cfd229ae3..5171d58bf 100644 --- a/src/rpucuda/cuda/pulsed_weight_updater.h +++ b/src/rpucuda/cuda/pulsed_weight_updater.h @@ -62,6 +62,17 @@ template class PulsedWeightUpdater { const bool x_trans, const bool d_trans, const T beta = (T)1.0); + + template + void doInfiniteGranularityUpdate( + XInputIteratorT x_in, + DInputIteratorT d_in, + PulsedRPUDeviceCuda *pulsed_device, + T *dev_weights, + const T lr, + const int m_batch, + const bool x_trans, + const bool d_trans); void setVerbosityLevel(int level) { verbose_ = level; }; void dumpExtra(RPU::state_t &extra, const std::string prefix); void loadExtra(const RPU::state_t &extra, const std::string prefix, bool strict); diff --git a/src/rpucuda/cuda/rpucuda_constantstep_device.cu b/src/rpucuda/cuda/rpucuda_constantstep_device.cu index 3c0623714..d6f9a0167 100644 --- a/src/rpucuda/cuda/rpucuda_constantstep_device.cu +++ b/src/rpucuda/cuda/rpucuda_constantstep_device.cu @@ -114,6 +114,54 @@ pwukpvec_t ConstantStepRPUDeviceCuda::getUpdateKernels( #undef ARGS +/*********************************************************************************/ +/* infinite granularity update */ + +namespace { +constexpr int IG_THREADS_PER_BLOCK = 256; + +inline int ig_num_blocks(int total_size) { + return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK; +} + +template +__global__ void kernelIGConstantStep( + T *weights, const T *grad_matrix, const param_t *params, int total_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_size) { + return; + } + + T G = grad_matrix[idx]; + if (G == (T)0.0) { + return; + } + + param4_t p = reinterpret_cast(params)[idx]; + T wmin = p.x; + T scale_down = p.y; + T wmax = p.z; + T scale_up = p.w; + + T response = (G < (T)0.0) ? scale_down : scale_up; + T w = weights[idx]; + w += G * response; + w = (w > wmax) ? wmax : w; + w = (w < wmin) ? wmin : w; + weights[idx] = w; +} +} // namespace + +template +void ConstantStepRPUDeviceCuda::doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + UNUSED(dev_states); + int total_size = this->x_size_ * this->d_size_; + kernelIGConstantStep + <<context_->getStream()>>>( + dev_weights, grad_matrix, this->get4ParamsData(), total_size); +} + template class ConstantStepRPUDeviceCuda; #ifdef RPU_USE_DOUBLE template class ConstantStepRPUDeviceCuda; diff --git a/src/rpucuda/cuda/rpucuda_constantstep_device.h b/src/rpucuda/cuda/rpucuda_constantstep_device.h index 7800709b0..6cb4314c3 100644 --- a/src/rpucuda/cuda/rpucuda_constantstep_device.h +++ b/src/rpucuda/cuda/rpucuda_constantstep_device.h @@ -55,6 +55,8 @@ template class ConstantStepRPUDeviceCuda : public PulsedRPUDeviceCu int use_bo64, bool out_trans, const PulsedUpdateMetaParameter &up) override; + void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override; }; } // namespace RPU diff --git a/src/rpucuda/cuda/rpucuda_expstep_device.cu b/src/rpucuda/cuda/rpucuda_expstep_device.cu index e80468e2c..1f4808f2f 100644 --- a/src/rpucuda/cuda/rpucuda_expstep_device.cu +++ b/src/rpucuda/cuda/rpucuda_expstep_device.cu @@ -207,6 +207,103 @@ pwukpvec_t ExpStepRPUDeviceCuda::getUpdateKernels( return v; } +/*********************************************************************************/ +/* infinite granularity update */ + +namespace { +constexpr int IG_THREADS_PER_BLOCK = 256; + +inline int ig_num_blocks(int total_size) { + return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK; +} + +template +__device__ __forceinline__ void igStoreWeight( + T *weights, + T *persistent_weights, + const T *write_noise_std, + curandState_t *random_states, + int idx, + T w) { + if (persistent_weights == nullptr) { + weights[idx] = w; + return; + } + + persistent_weights[idx] = w; + T uw_std = write_noise_std == nullptr ? (T)0.0 : *write_noise_std; + if (uw_std > (T)0.0 && random_states != nullptr) { + curandState local_state = random_states[idx]; + weights[idx] = w + uw_std * (T)curand_normal(&local_state); + random_states[idx] = local_state; + } else { + weights[idx] = w; + } +} + +template +__global__ void kernelIGExpStep( + T *weights, + const T *grad_matrix, + const param_t *params4, + const T *es_par, + int total_size, + T *persistent_weights, + curandState_t *random_states) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_size) { + return; + } + + T G = grad_matrix[idx]; + if (G == (T)0.0) { + return; + } + + param4_t p4 = reinterpret_cast(params4)[idx]; + T wmin = p4.x; + T wmax = p4.z; + T b_diff = wmax - wmin; + if (b_diff <= (T)0.0) { + return; + } + + T A_down = es_par[0]; + T A_up = es_par[1]; + T gamma_down = es_par[2]; + T gamma_up = es_par[3]; + T a = es_par[4]; + T b = es_par[5]; + + T w = persistent_weights == nullptr ? weights[idx] : persistent_weights[idx]; + T abs_G = (G > (T)0.0) ? G : -G; + T z = (T)2.0 * w / b_diff * a + b; + T dw; + if (G < (T)0.0) { + T y = (T)1.0 - A_down * __expf(-gamma_down * z); + dw = (y > (T)0.0) ? -y * (T)p4.y * abs_G : (T)0.0; + } else { + T y = (T)1.0 - A_up * __expf(gamma_up * z); + dw = (y > (T)0.0) ? y * (T)p4.w * abs_G : (T)0.0; + } + + w += dw; + w = (w > wmax) ? wmax : w; + w = (w < wmin) ? wmin : w; + igStoreWeight(weights, persistent_weights, es_par + 6, random_states, idx, w); +} +} // namespace + +template +void ExpStepRPUDeviceCuda::doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + int total_size = this->x_size_ * this->d_size_; + kernelIGExpStep + <<context_->getStream()>>>( + dev_weights, grad_matrix, this->get4ParamsData(), this->getGlobalParamsData(), + total_size, this->get1ParamsData(), dev_states); +} + template class ExpStepRPUDeviceCuda; #ifdef RPU_USE_DOUBLE template class ExpStepRPUDeviceCuda; diff --git a/src/rpucuda/cuda/rpucuda_expstep_device.h b/src/rpucuda/cuda/rpucuda_expstep_device.h index 117cc66a6..fbf72175a 100644 --- a/src/rpucuda/cuda/rpucuda_expstep_device.h +++ b/src/rpucuda/cuda/rpucuda_expstep_device.h @@ -42,7 +42,7 @@ template class ExpStepRPUDeviceCuda : public PulsedRPUDeviceCuda es_par_arr[3] = par.es_gamma_up; es_par_arr[4] = par.es_a; es_par_arr[5] = par.es_b; - es_par_arr[6] = par.getScaledWriteNoise(); + es_par_arr[6] = par.write_noise_std * (par.dw_min == (T)0.0 ? (T)1.0 : par.dw_min); es_par_arr[7] = par.dw_min_std_add; es_par_arr[8] = par.dw_min_std_slope; dev_es_par_->assign(es_par_arr); @@ -65,6 +65,9 @@ template class ExpStepRPUDeviceCuda : public PulsedRPUDeviceCuda : PulsedRPUDeviceCuda::getWeightGranularityNoise(); } + void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override; + private: std::unique_ptr> dev_es_par_ = nullptr; }; diff --git a/src/rpucuda/cuda/rpucuda_linearstep_device.cu b/src/rpucuda/cuda/rpucuda_linearstep_device.cu index d70c4e733..dfbbe942b 100644 --- a/src/rpucuda/cuda/rpucuda_linearstep_device.cu +++ b/src/rpucuda/cuda/rpucuda_linearstep_device.cu @@ -178,6 +178,90 @@ pwukpvec_t LinearStepRPUDeviceCuda::getUpdateKernels( #undef ARGS +/*********************************************************************************/ +/* infinite granularity update */ + +namespace { +constexpr int IG_THREADS_PER_BLOCK = 256; + +inline int ig_num_blocks(int total_size) { + return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK; +} + +template +__device__ __forceinline__ void igStoreWeight( + T *weights, + T *persistent_weights, + const T *write_noise_std, + curandState_t *random_states, + int idx, + T w) { + if (persistent_weights == nullptr) { + weights[idx] = w; + return; + } + + persistent_weights[idx] = w; + T uw_std = write_noise_std == nullptr ? (T)0.0 : *write_noise_std; + if (uw_std > (T)0.0 && random_states != nullptr) { + curandState local_state = random_states[idx]; + weights[idx] = w + uw_std * (T)curand_normal(&local_state); + random_states[idx] = local_state; + } else { + weights[idx] = w; + } +} + +template +__global__ void kernelIGLinearStep( + T *weights, + const T *grad_matrix, + const param_t *params4, + const param_t *params2, + int total_size, + T *persistent_weights, + const T *write_noise_std, + curandState_t *random_states) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_size) { + return; + } + + T G = grad_matrix[idx]; + if (G == (T)0.0) { + return; + } + + param4_t p4 = reinterpret_cast(params4)[idx]; + param2_t p2 = reinterpret_cast(params2)[idx]; + + T wmin = p4.x; + T scale_down = p4.y; + T wmax = p4.z; + T scale_up = p4.w; + T slope_down = p2.x; + T slope_up = p2.y; + + T w = persistent_weights == nullptr ? weights[idx] : persistent_weights[idx]; + T response = G < (T)0.0 ? slope_down * w + scale_down : slope_up * w + scale_up; + + w += G * response; + w = (w > wmax) ? wmax : w; + w = (w < wmin) ? wmin : w; + igStoreWeight(weights, persistent_weights, write_noise_std, random_states, idx, w); +} +} // namespace + +template +void LinearStepRPUDeviceCuda::doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + int total_size = this->x_size_ * this->d_size_; + kernelIGLinearStep + <<context_->getStream()>>>( + dev_weights, grad_matrix, this->get4ParamsData(), dev_slope_->getDataConst(), total_size, + this->get1ParamsData(), this->getGlobalParamsData(), dev_states); +} + template class LinearStepRPUDeviceCuda; #ifdef RPU_USE_DOUBLE template class LinearStepRPUDeviceCuda; diff --git a/src/rpucuda/cuda/rpucuda_linearstep_device.h b/src/rpucuda/cuda/rpucuda_linearstep_device.h index 5b545de1d..c2f34fd71 100644 --- a/src/rpucuda/cuda/rpucuda_linearstep_device.h +++ b/src/rpucuda/cuda/rpucuda_linearstep_device.h @@ -51,7 +51,8 @@ template class LinearStepRPUDeviceCuda : public PulsedRPUDeviceCuda tmp_slope[kk + 1] = w_slope_up[i][j]; } } dev_slope_->assign(tmp_slope); - dev_write_noise_std_->setConst(getPar().getScaledWriteNoise()); + dev_write_noise_std_->setConst( + getPar().write_noise_std * (getPar().dw_min == (T)0.0 ? (T)1.0 : getPar().dw_min)); this->context_->synchronize(); delete[] tmp_slope;); @@ -67,6 +68,8 @@ template class LinearStepRPUDeviceCuda : public PulsedRPUDeviceCuda T *get1ParamsData() override { return getPar().usesPersistentWeight() ? this->dev_persistent_weights_->getData() : nullptr; }; + void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override; T getWeightGranularityNoise() const override { // need to make sure that random states are enabled return getPar().usesPersistentWeight() diff --git a/src/rpucuda/cuda/rpucuda_piecewisestep_device.cu b/src/rpucuda/cuda/rpucuda_piecewisestep_device.cu index 38849dd54..7944f1b9f 100644 --- a/src/rpucuda/cuda/rpucuda_piecewisestep_device.cu +++ b/src/rpucuda/cuda/rpucuda_piecewisestep_device.cu @@ -12,7 +12,6 @@ namespace RPU { namespace { - template __device__ __forceinline__ T get_interpolated_scale( const T &w, @@ -155,6 +154,90 @@ pwukpvec_t PiecewiseStepRPUDeviceCuda::getUpdateKernels( #undef ARGS #undef ADD_KERNELS +/*********************************************************************************/ +/* infinite granularity update */ + +namespace { +constexpr int IG_THREADS_PER_BLOCK = 256; + +inline int ig_num_blocks(int total_size) { + return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK; +} + +template +__device__ __forceinline__ void igStoreWeight( + T *weights, + T *persistent_weights, + const T *write_noise_std, + curandState_t *random_states, + int idx, + T w) { + if (persistent_weights == nullptr) { + weights[idx] = w; + return; + } + + persistent_weights[idx] = w; + T uw_std = write_noise_std == nullptr ? (T)0.0 : *write_noise_std; + if (uw_std > (T)0.0 && random_states != nullptr) { + curandState local_state = random_states[idx]; + weights[idx] = w + uw_std * (T)curand_normal(&local_state); + random_states[idx] = local_state; + } else { + weights[idx] = w; + } +} + +template +__global__ void kernelIGPiecewiseStep( + T *weights, + const T *grad_matrix, + const param_t *params4, + const T *global_pars, + int gp_count, + int total_size, + T *persistent_weights, + curandState_t *random_states) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_size) { + return; + } + + T G = grad_matrix[idx]; + if (G == (T)0.0) { + return; + } + + param4_t p4 = reinterpret_cast(params4)[idx]; + T wmin = p4.x; + T wmax = p4.z; + T w_range = wmax - wmin; + + T abs_G = (G > (T)0.0) ? G : -G; + T w = persistent_weights == nullptr ? weights[idx] : persistent_weights[idx]; + const T *piecewise_vec = + (G > (T)0.0) ? global_pars : global_pars + (size_t)gp_count / 2; + const int n_points = (int)round(global_pars[gp_count / 2 - 1]); + const int n_sections = n_points - 1; + T scale = (G > (T)0.0) ? (T)p4.w * abs_G : -(T)p4.y * abs_G; + + w += get_interpolated_scale(w, piecewise_vec, scale, n_sections, w_range, wmin); + w = (w > wmax) ? wmax : w; + w = (w < wmin) ? wmin : w; + igStoreWeight(weights, persistent_weights, global_pars + gp_count - 1, random_states, idx, w); +} +} // namespace + +template +void PiecewiseStepRPUDeviceCuda::doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + int total_size = this->x_size_ * this->d_size_; + kernelIGPiecewiseStep + <<context_->getStream()>>>( + dev_weights, grad_matrix, this->get4ParamsData(), this->getGlobalParamsData(), gp_count_, + total_size, this->get1ParamsData(), dev_states); +} + template class PiecewiseStepRPUDeviceCuda; #ifdef RPU_USE_DOUBLE template class PiecewiseStepRPUDeviceCuda; diff --git a/src/rpucuda/cuda/rpucuda_piecewisestep_device.h b/src/rpucuda/cuda/rpucuda_piecewisestep_device.h index 8a73070ba..264a7072c 100644 --- a/src/rpucuda/cuda/rpucuda_piecewisestep_device.h +++ b/src/rpucuda/cuda/rpucuda_piecewisestep_device.h @@ -52,7 +52,8 @@ template class PiecewiseStepRPUDeviceCuda : public PulsedRPUDeviceC // n_points info tmp_global_pars[gp_count_ / 2 - 1] = (T)n_points; // write std info - tmp_global_pars[gp_count_ - 1] = getPar().getScaledWriteNoise(); + tmp_global_pars[gp_count_ - 1] = + getPar().write_noise_std * (getPar().dw_min == (T)0.0 ? (T)1.0 : getPar().dw_min); dev_global_pars_ = RPU::make_unique>(this->context_, gp_count_, tmp_global_pars); this->context_->synchronize(); @@ -76,6 +77,9 @@ template class PiecewiseStepRPUDeviceCuda : public PulsedRPUDeviceC : PulsedRPUDeviceCuda::getWeightGranularityNoise(); } + void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override; + private: std::unique_ptr> dev_global_pars_ = nullptr; int gp_count_ = 0; diff --git a/src/rpucuda/cuda/rpucuda_powstep_device.cu b/src/rpucuda/cuda/rpucuda_powstep_device.cu index 638df386e..b575e2ea4 100644 --- a/src/rpucuda/cuda/rpucuda_powstep_device.cu +++ b/src/rpucuda/cuda/rpucuda_powstep_device.cu @@ -111,6 +111,93 @@ pwukpvec_t PowStepRPUDeviceCuda::getUpdateKernels( #undef ARGS +/*********************************************************************************/ +/* infinite granularity update */ + +namespace { +constexpr int IG_THREADS_PER_BLOCK = 256; + +inline int ig_num_blocks(int total_size) { + return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK; +} + +template +__device__ __forceinline__ void igStoreWeight( + T *weights, + T *persistent_weights, + const T *write_noise_std, + curandState_t *random_states, + int idx, + T w) { + if (persistent_weights == nullptr) { + weights[idx] = w; + return; + } + + persistent_weights[idx] = w; + T uw_std = write_noise_std == nullptr ? (T)0.0 : *write_noise_std; + if (uw_std > (T)0.0 && random_states != nullptr) { + curandState local_state = random_states[idx]; + weights[idx] = w + uw_std * (T)curand_normal(&local_state); + random_states[idx] = local_state; + } else { + weights[idx] = w; + } +} + +template +__global__ void kernelIGPowStep( + T *weights, + const T *grad_matrix, + const param_t *params4, + const param_t *params_gamma, + int total_size, + T *persistent_weights, + const T *write_noise_std, + curandState_t *random_states) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_size) { + return; + } + + T G = grad_matrix[idx]; + if (G == (T)0.0) { + return; + } + + param4_t p4 = reinterpret_cast(params4)[idx]; + param2_t pg = reinterpret_cast(params_gamma)[idx]; + + T wmin = p4.x; + T wmax = p4.z; + T range = wmax - wmin; + if (range == (T)0.0) { + return; + } + + T abs_G = (G > (T)0.0) ? G : -G; + T w = persistent_weights == nullptr ? weights[idx] : persistent_weights[idx]; + T x = (wmax - w) / range; + T dw = G < (T)0.0 ? -(T)p4.y * abs_G * __powf((T)1.0 - x, (T)pg.x) + : (T)p4.w * abs_G * __powf(x, (T)pg.y); + + w += dw; + w = (w > wmax) ? wmax : w; + w = (w < wmin) ? wmin : w; + igStoreWeight(weights, persistent_weights, write_noise_std, random_states, idx, w); +} +} // namespace + +template +void PowStepRPUDeviceCuda::doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + int total_size = this->x_size_ * this->d_size_; + kernelIGPowStep + <<context_->getStream()>>>( + dev_weights, grad_matrix, this->get4ParamsData(), this->get2ParamsData(), total_size, + this->get1ParamsData(), this->getGlobalParamsData(), dev_states); +} + template class PowStepRPUDeviceCuda; #ifdef RPU_USE_DOUBLE template class PowStepRPUDeviceCuda; diff --git a/src/rpucuda/cuda/rpucuda_powstep_device.h b/src/rpucuda/cuda/rpucuda_powstep_device.h index a59284845..3b007e266 100644 --- a/src/rpucuda/cuda/rpucuda_powstep_device.h +++ b/src/rpucuda/cuda/rpucuda_powstep_device.h @@ -51,7 +51,8 @@ template class PowStepRPUDeviceCuda : public PulsedRPUDeviceCuda tmp_gamma[kk + 1] = w_gamma_up[i][j]; } } dev_gamma_->assign(tmp_gamma); - dev_write_noise_std_->setConst(getPar().getScaledWriteNoise()); + dev_write_noise_std_->setConst( + getPar().write_noise_std * (getPar().dw_min == (T)0.0 ? (T)1.0 : getPar().dw_min)); this->context_->synchronize(); delete[] tmp_gamma;); @@ -74,6 +75,9 @@ template class PowStepRPUDeviceCuda : public PulsedRPUDeviceCuda : PulsedRPUDeviceCuda::getWeightGranularityNoise(); } + void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override; + private: std::unique_ptr> dev_gamma_ = nullptr; std::unique_ptr> dev_write_noise_std_ = nullptr; diff --git a/src/rpucuda/cuda/rpucuda_powstep_reference_device.cu b/src/rpucuda/cuda/rpucuda_powstep_reference_device.cu index 8de88c8bd..a7942d294 100644 --- a/src/rpucuda/cuda/rpucuda_powstep_reference_device.cu +++ b/src/rpucuda/cuda/rpucuda_powstep_reference_device.cu @@ -118,6 +118,69 @@ pwukpvec_t PowStepReferenceRPUDeviceCuda::getUpdateKernels( #undef ARGS +/*********************************************************************************/ +/* infinite granularity update */ + +namespace { +constexpr int IG_THREADS_PER_BLOCK = 256; + +inline int ig_num_blocks(int total_size) { + return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK; +} + +template +__global__ void kernelIGPowStepReference( + T *weights, + const T *grad_matrix, + const param_t *params4, + const param_t *params_gamma, + const T *reference, + int total_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_size) { + return; + } + + T G = grad_matrix[idx]; + if (G == (T)0.0) { + return; + } + + param4_t p4 = reinterpret_cast(params4)[idx]; + param2_t pg = reinterpret_cast(params_gamma)[idx]; + + T wmin = p4.x; + T wmax = p4.z; + T range = wmax - wmin; + if (range == (T)0.0) { + return; + } + + T ref = reference[idx]; + T abs_G = (G > (T)0.0) ? G : -G; + T w = weights[idx] + ref; + T x = (wmax - w) / range; + T dw = G < (T)0.0 ? -(T)p4.y * abs_G * __powf((T)1.0 - x, (T)pg.x) + : (T)p4.w * abs_G * __powf(x, (T)pg.y); + + w += dw; + w = (w > wmax) ? wmax : w; + w = (w < wmin) ? wmin : w; + weights[idx] = w - ref; +} +} // namespace + +template +void PowStepReferenceRPUDeviceCuda::doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + UNUSED(dev_states); + int total_size = this->x_size_ * this->d_size_; + kernelIGPowStepReference + <<context_->getStream()>>>( + dev_weights, grad_matrix, this->get4ParamsData(), this->get2ParamsData(), + this->get1ParamsData(), total_size); +} + template class PowStepReferenceRPUDeviceCuda; #ifdef RPU_USE_DOUBLE template class PowStepReferenceRPUDeviceCuda; diff --git a/src/rpucuda/cuda/rpucuda_powstep_reference_device.h b/src/rpucuda/cuda/rpucuda_powstep_reference_device.h index 6aeab6ee6..2ebf8148a 100644 --- a/src/rpucuda/cuda/rpucuda_powstep_reference_device.h +++ b/src/rpucuda/cuda/rpucuda_powstep_reference_device.h @@ -73,6 +73,9 @@ template class PowStepReferenceRPUDeviceCuda : public PulsedRPUDevi param_t *get2ParamsData() override { return dev_gamma_->getData(); }; T *get1ParamsData() override { return dev_reference_->getData(); }; + void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override; + private: std::unique_ptr> dev_gamma_ = nullptr; std::unique_ptr> dev_reference_ = nullptr; diff --git a/src/rpucuda/cuda/rpucuda_pulsed_device.cu b/src/rpucuda/cuda/rpucuda_pulsed_device.cu index f41eca18a..3ea8340f5 100644 --- a/src/rpucuda/cuda/rpucuda_pulsed_device.cu +++ b/src/rpucuda/cuda/rpucuda_pulsed_device.cu @@ -495,6 +495,18 @@ void PulsedRPUDeviceCuda::runUpdateKernel( d_counts_chunk, cwo); } +/*********************************************************************************/ +/* infinite granularity update — default (ConstantStep) */ + +template +void PulsedRPUDeviceCuda::doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + UNUSED(dev_weights); + UNUSED(grad_matrix); + UNUSED(dev_states); + RPU_FATAL("Infinite granularity update not available for this device!"); +} + template class PulsedRPUDeviceCuda; #ifdef RPU_USE_DOUBLE template class PulsedRPUDeviceCuda; diff --git a/src/rpucuda/cuda/rpucuda_pulsed_device.h b/src/rpucuda/cuda/rpucuda_pulsed_device.h index f2c0c43d4..cf4eed0a3 100644 --- a/src/rpucuda/cuda/rpucuda_pulsed_device.h +++ b/src/rpucuda/cuda/rpucuda_pulsed_device.h @@ -68,6 +68,10 @@ template class PulsedRPUDeviceCudaBase : public SimpleRPUDeviceCuda bool isPulsedDevice() const override { return true; }; bool hasDirectUpdate() const override { return false; }; + virtual void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + RPU_FATAL("Infinite granularity update not available for this device!"); + }; void doDirectUpdate( const T *x_input, const T *d_input, @@ -167,6 +171,8 @@ template class PulsedRPUDeviceCuda : public PulsedRPUDeviceCudaBase void resetCols(T *dev_weights, int start_col, int n_cols, T reset_prob) override; virtual void resetAt(T *dev_weights, const char *dev_non_zero_msk); void applyWeightUpdate(T *dev_weights, T *dw_and_current_weight_out) override; + void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override; void populateFrom(const AbstractRPUDevice &rpu_device) override; // need to be called by derived PulsedRPUDeviceCuda *clone() const override { RPU_FATAL("Needs implementations"); }; diff --git a/src/rpucuda/cuda/rpucuda_softbounds_reference_device.cu b/src/rpucuda/cuda/rpucuda_softbounds_reference_device.cu index 44410ca20..d4dee0ba3 100644 --- a/src/rpucuda/cuda/rpucuda_softbounds_reference_device.cu +++ b/src/rpucuda/cuda/rpucuda_softbounds_reference_device.cu @@ -211,6 +211,100 @@ pwukpvec_t SoftBoundsReferenceRPUDeviceCuda::getUpdateKernels( #undef ARGS +/*********************************************************************************/ +/* infinite granularity update */ + +namespace { +constexpr int IG_THREADS_PER_BLOCK = 256; + +inline int ig_num_blocks(int total_size) { + return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK; +} + +template +__device__ __forceinline__ void igStoreWeight( + T *weights, + T *persistent_weights, + const T *write_noise_std, + curandState_t *random_states, + int idx, + T w) { + if (persistent_weights == nullptr) { + weights[idx] = w; + return; + } + + persistent_weights[idx] = w; + T uw_std = write_noise_std == nullptr ? (T)0.0 : *write_noise_std; + if (uw_std > (T)0.0 && random_states != nullptr) { + curandState local_state = random_states[idx]; + weights[idx] = w + uw_std * (T)curand_normal(&local_state); + random_states[idx] = local_state; + } else { + weights[idx] = w; + } +} + +template +__global__ void kernelIGSoftBoundsReference( + T *weights, + const T *grad_matrix, + const param_t *params4, + const param_t *params_ref, + int total_size, + T *persistent_weights, + const T *write_noise_std, + curandState_t *random_states) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_size) { + return; + } + + T G = grad_matrix[idx]; + if (G == (T)0.0) { + return; + } + + param4_t p4 = reinterpret_cast(params4)[idx]; + param2_t pr = reinterpret_cast(params_ref)[idx]; + + T wmin = p4.x; + T wmax = p4.z; + T ref = pr.x; + T abs_G = (G > (T)0.0) ? G : -G; + T w_shifted = (persistent_weights == nullptr ? weights[idx] : persistent_weights[idx]) + ref; + + T lin_a = (T)0.0; + T lin_dw = (T)0.0; + if (G > (T)0.0) { + lin_dw = (T)p4.w * abs_G; + if (wmax > (T)0.0) { + lin_a = -lin_dw / wmax; + } + } else { + lin_dw = -(T)p4.y * abs_G; + if (wmin < (T)0.0) { + lin_a = -lin_dw / wmin; + } + } + + w_shifted += lin_a * w_shifted + lin_dw; + w_shifted = (w_shifted > wmax) ? wmax : w_shifted; + w_shifted = (w_shifted < wmin) ? wmin : w_shifted; + igStoreWeight(weights, persistent_weights, write_noise_std, random_states, idx, w_shifted - ref); +} +} // namespace + +template +void SoftBoundsReferenceRPUDeviceCuda::doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) { + int total_size = this->x_size_ * this->d_size_; + kernelIGSoftBoundsReference + <<context_->getStream()>>>( + dev_weights, grad_matrix, this->get4ParamsData(), this->get2ParamsData(), total_size, + this->get1ParamsData(), this->getGlobalParamsData(), dev_states); +} + template class SoftBoundsReferenceRPUDeviceCuda; #ifdef RPU_USE_DOUBLE template class SoftBoundsReferenceRPUDeviceCuda; diff --git a/src/rpucuda/cuda/rpucuda_softbounds_reference_device.h b/src/rpucuda/cuda/rpucuda_softbounds_reference_device.h index 257c68284..8268694c7 100644 --- a/src/rpucuda/cuda/rpucuda_softbounds_reference_device.h +++ b/src/rpucuda/cuda/rpucuda_softbounds_reference_device.h @@ -52,7 +52,8 @@ template class SoftBoundsReferenceRPUDeviceCuda : public PulsedRPUD } dev_reference_->assign(tmp); - dev_write_noise_std_->setConst(getPar().getScaledWriteNoise()); + dev_write_noise_std_->setConst( + getPar().write_noise_std * (getPar().dw_min == (T)0.0 ? (T)1.0 : getPar().dw_min)); this->context_->synchronize(); delete[] tmp;); @@ -75,6 +76,9 @@ template class SoftBoundsReferenceRPUDeviceCuda : public PulsedRPUD : PulsedRPUDeviceCuda::getWeightGranularityNoise(); } + void doInfiniteGranularityUpdate( + T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override; + private: std::unique_ptr> dev_reference_ = nullptr; std::unique_ptr> dev_write_noise_std_ = nullptr; diff --git a/src/rpucuda/rpu_constantstep_device.cpp b/src/rpucuda/rpu_constantstep_device.cpp index d284b35c7..f99c0a0d0 100644 --- a/src/rpucuda/rpu_constantstep_device.cpp +++ b/src/rpucuda/rpu_constantstep_device.cpp @@ -77,6 +77,19 @@ void ConstantStepRPUDevice::doDenseUpdate(T **weights, int *coincidences, RNG ); } +template +void ConstantStepRPUDevice::doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + PulsedRPUDevice::doInfiniteGranularityUpdate( + weights, x_input, x_inc, d_input, d_inc, learning_rate, rng); +} + template class ConstantStepRPUDevice; #ifdef RPU_USE_DOUBLE template class ConstantStepRPUDevice; diff --git a/src/rpucuda/rpu_constantstep_device.h b/src/rpucuda/rpu_constantstep_device.h index 55e2c6068..15294f3e9 100644 --- a/src/rpucuda/rpu_constantstep_device.h +++ b/src/rpucuda/rpu_constantstep_device.h @@ -55,5 +55,14 @@ template class ConstantStepRPUDevice : public PulsedRPUDevice { T **weights, int i, const int *x_signed_indices, int x_count, int d_sign, RNG *rng) override; void doDenseUpdate(T **weights, int *coincidences, RNG *rng) override; + void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) + override; }; } // namespace RPU diff --git a/src/rpucuda/rpu_expstep_device.cpp b/src/rpucuda/rpu_expstep_device.cpp index 7293bb80e..1c21caee8 100644 --- a/src/rpucuda/rpu_expstep_device.cpp +++ b/src/rpucuda/rpu_expstep_device.cpp @@ -13,7 +13,7 @@ void ExpStepRPUDevice::populate( const ExpStepRPUDeviceMetaParameter &p, RealWorldRNG *rng) { PulsedRPUDevice::populate(p, rng); } -template +template inline void update_once( T &w, T &w_apparent, @@ -30,7 +30,7 @@ inline void update_once( const T &es_gamma_up, const T &dw_min_std, const T &write_noise_std, - RNG *rng) { + RngT *rng) { T b_diff = (max_bound - min_bound); if (b_diff > (T)0.0) { T z = (T)2.0 * w / b_diff * es_a + es_b; @@ -54,7 +54,7 @@ inline void update_once( } } -template +template inline void update_once_complex_noise( T &w, T &w_apparent, @@ -73,7 +73,7 @@ inline void update_once_complex_noise( const T &dw_min_std_add, const T &dw_min_std_slope, const T &write_noise_std, - RNG *rng) { + RngT *rng) { T b_diff = (max_bound - min_bound); if (b_diff > (T)0.0) { T z = (T)2.0 * w / b_diff * es_a + es_b; @@ -159,6 +159,45 @@ void ExpStepRPUDevice::doDenseUpdate(T **weights, int *coincidences, RNG * } } +/* infinite granularity update */ +template +void ExpStepRPUDevice::doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + + const auto &par = getPar(); + T write_noise_std = par.write_noise_std * (par.dw_min == (T)0.0 ? (T)1.0 : par.dw_min); + + for (int i = 0; i < this->d_size_; i++) { + T d_val = d_input[i * d_inc]; + if (d_val == (T)0.0) { + continue; + } + + T *scale_down = this->w_scale_down_[i]; + T *scale_up = this->w_scale_up_[i]; + T *w = par.usesPersistentWeight() ? this->w_persistent_[i] : weights[i]; + T *w_apparent = weights[i]; + T *min_bound = this->w_min_bound_[i]; + T *max_bound = this->w_max_bound_[i]; + + IG_UPDATE_W_LOOP_INNER( + T sd = scale_down[j] * abs_G_lr; + T su = scale_up[j] * abs_G_lr; + update_once( + w[j], w_apparent[j], sign, min_bound[j], max_bound[j], sd, su, + par.es_a, par.es_b, par.es_A_down, par.es_A_up, + par.es_gamma_down, par.es_gamma_up, (T)0.0, write_noise_std, + rng); + ); + } +} + template class ExpStepRPUDevice; #ifdef RPU_USE_DOUBLE diff --git a/src/rpucuda/rpu_expstep_device.h b/src/rpucuda/rpu_expstep_device.h index 623f374fd..1053fbd34 100644 --- a/src/rpucuda/rpu_expstep_device.h +++ b/src/rpucuda/rpu_expstep_device.h @@ -87,6 +87,15 @@ template class ExpStepRPUDevice : public PulsedRPUDevice { T **weights, int i, const int *x_signed_indices, int x_count, int d_sign, RNG *rng) override; void doDenseUpdate(T **weights, int *coincidences, RNG *rng) override; + void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) + override; }; } // namespace RPU diff --git a/src/rpucuda/rpu_linearstep_device.cpp b/src/rpucuda/rpu_linearstep_device.cpp index c730d5c0f..2beca309d 100644 --- a/src/rpucuda/rpu_linearstep_device.cpp +++ b/src/rpucuda/rpu_linearstep_device.cpp @@ -221,6 +221,50 @@ void LinearStepRPUDevice::doDenseUpdate(T **weights, int *coincidences, RNG +void LinearStepRPUDevice::doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + + const auto &par = getPar(); + T write_noise_std = par.write_noise_std * (par.dw_min == (T)0.0 ? (T)1.0 : par.dw_min); + + for (int i = 0; i < this->d_size_; i++) { + T d_val = d_input[i * d_inc]; + if (d_val == (T)0.0) { + continue; + } + + T *scale_down = this->w_scale_down_[i]; + T *scale_up = this->w_scale_up_[i]; + T *slope_down = w_slope_down_[i]; + T *slope_up = w_slope_up_[i]; + T *w = par.usesPersistentWeight() ? this->w_persistent_[i] : weights[i]; + T *w_apparent = weights[i]; + T *min_bound = this->w_min_bound_[i]; + T *max_bound = this->w_max_bound_[i]; + + IG_UPDATE_W_LOOP_INNER( + T sd = scale_down[j] * abs_G_lr; + T su = scale_up[j] * abs_G_lr; + T sld = slope_down[j] * abs_G_lr; + T slu = slope_up[j] * abs_G_lr; + update_once_mult( + w[j], w_apparent[j], sign, sd, su, sld, slu, + min_bound[j], max_bound[j], (T)0.0, write_noise_std, + rng); + ); + } +} + template class LinearStepRPUDevice; #ifdef RPU_USE_DOUBLE template class LinearStepRPUDevice; diff --git a/src/rpucuda/rpu_linearstep_device.h b/src/rpucuda/rpu_linearstep_device.h index 0b5e735d2..f79c37f07 100644 --- a/src/rpucuda/rpu_linearstep_device.h +++ b/src/rpucuda/rpu_linearstep_device.h @@ -195,6 +195,15 @@ template class LinearStepRPUDevice : public PulsedRPUDevice { T **weights, int i, const int *x_signed_indices, int x_count, int d_sign, RNG *rng) override; void doDenseUpdate(T **weights, int *coincidences, RNG *rng) override; + void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) + override; private: T **w_slope_up_ = nullptr; diff --git a/src/rpucuda/rpu_piecewisestep_device.cpp b/src/rpucuda/rpu_piecewisestep_device.cpp index 64948181e..43cfa6207 100644 --- a/src/rpucuda/rpu_piecewisestep_device.cpp +++ b/src/rpucuda/rpu_piecewisestep_device.cpp @@ -25,7 +25,7 @@ void PiecewiseStepRPUDevice::populate( } namespace { -template +template inline void update_once( T &w, T &w_apparent, @@ -38,7 +38,7 @@ inline void update_once( const std::vector &piecewise_down_vec, const T &dw_min_std, const T &write_noise_std, - RNG *rng) { + RngT *rng) { size_t n_sections = piecewise_up_vec.size() - 1; @@ -120,6 +120,44 @@ void PiecewiseStepRPUDevice::doDenseUpdate(T **weights, int *coincidences, RN par.piecewise_down_vec, par.dw_min_std, write_noise_std, rng);); } +/* infinite granularity update */ +template +void PiecewiseStepRPUDevice::doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + + const auto &par = getPar(); + T write_noise_std = par.write_noise_std * (par.dw_min == (T)0.0 ? (T)1.0 : par.dw_min); + + for (int i = 0; i < this->d_size_; i++) { + T d_val = d_input[i * d_inc]; + if (d_val == (T)0.0) { + continue; + } + + T *scale_down = this->w_scale_down_[i]; + T *scale_up = this->w_scale_up_[i]; + T *w = par.usesPersistentWeight() ? this->w_persistent_[i] : weights[i]; + T *w_apparent = weights[i]; + T *min_bound = this->w_min_bound_[i]; + T *max_bound = this->w_max_bound_[i]; + + IG_UPDATE_W_LOOP_INNER( + T sd = scale_down[j] * abs_G_lr; + T su = scale_up[j] * abs_G_lr; + update_once( + w[j], w_apparent[j], sign, sd, su, + min_bound[j], max_bound[j], par.piecewise_up_vec, par.piecewise_down_vec, + (T)0.0, write_noise_std, rng); + ); + } +} + template class PiecewiseStepRPUDevice; #ifdef RPU_USE_DOUBLE template class PiecewiseStepRPUDevice; diff --git a/src/rpucuda/rpu_piecewisestep_device.h b/src/rpucuda/rpu_piecewisestep_device.h index 6027d830b..b207ffd65 100644 --- a/src/rpucuda/rpu_piecewisestep_device.h +++ b/src/rpucuda/rpu_piecewisestep_device.h @@ -75,5 +75,14 @@ template class PiecewiseStepRPUDevice : public PulsedRPUDevice { T **weights, int i, const int *x_signed_indices, int x_count, int d_sign, RNG *rng) override; void doDenseUpdate(T **weights, int *coincidences, RNG *rng) override; + void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) + override; }; } // namespace RPU diff --git a/src/rpucuda/rpu_powstep_device.cpp b/src/rpucuda/rpu_powstep_device.cpp index f4572ef81..1dd7c2d7a 100644 --- a/src/rpucuda/rpu_powstep_device.cpp +++ b/src/rpucuda/rpu_powstep_device.cpp @@ -85,7 +85,7 @@ template void PowStepRPUDevice::printDP(int x_count, int d_count } namespace { -template +template inline void update_once( T &w, T &w_apparent, @@ -98,7 +98,7 @@ inline void update_once( T &max_bound, const T &dw_min_std, const T &write_noise_std, - RNG *rng) { + RngT *rng) { T range = max_bound - min_bound; if (range == (T)0.0) { return; @@ -163,6 +163,46 @@ void PowStepRPUDevice::doDenseUpdate(T **weights, int *coincidences, RNG * par.dw_min_std, write_noise_std, rng);); } +/* infinite granularity update */ +template +void PowStepRPUDevice::doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + + const auto &par = getPar(); + T write_noise_std = par.write_noise_std * (par.dw_min == (T)0.0 ? (T)1.0 : par.dw_min); + + for (int i = 0; i < this->d_size_; i++) { + T d_val = d_input[i * d_inc]; + if (d_val == (T)0.0) { + continue; + } + + T *scale_down = this->w_scale_down_[i]; + T *scale_up = this->w_scale_up_[i]; + T *gamma_down = w_gamma_down_[i]; + T *gamma_up = w_gamma_up_[i]; + T *w = par.usesPersistentWeight() ? this->w_persistent_[i] : weights[i]; + T *w_apparent = weights[i]; + T *min_bound = this->w_min_bound_[i]; + T *max_bound = this->w_max_bound_[i]; + + IG_UPDATE_W_LOOP_INNER( + T sd = scale_down[j] * abs_G_lr; + T su = scale_up[j] * abs_G_lr; + update_once( + w[j], w_apparent[j], sign, sd, su, gamma_down[j], gamma_up[j], + min_bound[j], max_bound[j], (T)0.0, write_noise_std, + rng); + ); + } +} + template class PowStepRPUDevice; #ifdef RPU_USE_DOUBLE template class PowStepRPUDevice; diff --git a/src/rpucuda/rpu_powstep_device.h b/src/rpucuda/rpu_powstep_device.h index 0cb1f9e38..cf3ccf0fa 100644 --- a/src/rpucuda/rpu_powstep_device.h +++ b/src/rpucuda/rpu_powstep_device.h @@ -115,6 +115,15 @@ template class PowStepRPUDevice : public PulsedRPUDevice { T **weights, int i, const int *x_signed_indices, int x_count, int d_sign, RNG *rng) override; void doDenseUpdate(T **weights, int *coincidences, RNG *rng) override; + void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) + override; private: T **w_gamma_up_ = nullptr; diff --git a/src/rpucuda/rpu_powstep_reference_device.cpp b/src/rpucuda/rpu_powstep_reference_device.cpp index 2d5fec6f2..2b8353664 100644 --- a/src/rpucuda/rpu_powstep_reference_device.cpp +++ b/src/rpucuda/rpu_powstep_reference_device.cpp @@ -114,7 +114,7 @@ template void PowStepReferenceRPUDevice::printDP(int x_count, in } namespace { -template +template inline void update_once_reference( T &w, int &sign, @@ -126,7 +126,7 @@ inline void update_once_reference( T &min_bound, T &max_bound, const T &dw_min_std, - RNG *rng) { + RngT *rng) { T range = max_bound - min_bound; if (range == (T)0.0) { return; @@ -187,6 +187,42 @@ void PowStepReferenceRPUDevice::doDenseUpdate(T **weights, int *coincidences, ref[j], min_bound[j], max_bound[j], par.dw_min_std, rng);); } +/* infinite granularity update */ +template +void PowStepReferenceRPUDevice::doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + + for (int i = 0; i < this->d_size_; i++) { + T d_val = d_input[i * d_inc]; + if (d_val == (T)0.0) { + continue; + } + + T *scale_down = this->w_scale_down_[i]; + T *scale_up = this->w_scale_up_[i]; + T *gamma_down = w_gamma_down_[i]; + T *gamma_up = w_gamma_up_[i]; + T *ref = w_reference_[i]; + T *w = weights[i]; + T *min_bound = this->w_min_bound_[i]; + T *max_bound = this->w_max_bound_[i]; + + IG_UPDATE_W_LOOP_INNER( + T sd = scale_down[j] * abs_G_lr; + T su = scale_up[j] * abs_G_lr; + update_once_reference( + w[j], sign, sd, su, gamma_down[j], gamma_up[j], + ref[j], min_bound[j], max_bound[j], (T)0.0, rng); + ); + } +} + template class PowStepReferenceRPUDevice; #ifdef RPU_USE_DOUBLE template class PowStepReferenceRPUDevice; diff --git a/src/rpucuda/rpu_powstep_reference_device.h b/src/rpucuda/rpu_powstep_reference_device.h index ba171b938..ee645a6e5 100644 --- a/src/rpucuda/rpu_powstep_reference_device.h +++ b/src/rpucuda/rpu_powstep_reference_device.h @@ -143,6 +143,15 @@ template class PowStepReferenceRPUDevice : public PulsedRPUDevice *rng) override; void doDenseUpdate(T **weights, int *coincidences, RNG *rng) override; + void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) + override; private: T **w_gamma_up_ = nullptr; diff --git a/src/rpucuda/rpu_pulsed_device.cpp b/src/rpucuda/rpu_pulsed_device.cpp index bb4ae5c73..cf5a06769 100644 --- a/src/rpucuda/rpu_pulsed_device.cpp +++ b/src/rpucuda/rpu_pulsed_device.cpp @@ -587,6 +587,40 @@ template void PulsedRPUDevice::applyUpdateWriteNoise(T **weights } } +/*********************************************************************************/ +/* infinite granularity update (dw_min=0 mean-field limit) */ + +template +void PulsedRPUDevice::doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + UNUSED(rng); + // Default implementation: ConstantStep response (no weight dependence). + // Devices with weight-dependent response (LinearStep, ExpStep, PowStep, etc.) + // override this method. + for (int i = 0; i < this->d_size_; i++) { + T d_val = d_input[i * d_inc]; + if (d_val == (T)0.0) { + continue; + } + T *w = weights[i]; + T *max_bound = w_max_bound_[i]; + T *min_bound = w_min_bound_[i]; + IG_UPDATE_W_LOOP_INNER( + T sd = w_scale_down_[i][j] * abs_G_lr; + T su = w_scale_up_[i][j] * abs_G_lr; + w[j] += (sign > 0) ? -sd : su; + w[j] = MIN(w[j], max_bound[j]); + w[j] = MAX(w[j], min_bound[j]); + ); + } +} + /*********************************************************************************/ /* populate */ @@ -603,7 +637,14 @@ void PulsedRPUDevice::populate(const PulsedRPUDeviceMetaParameter &p, Real T up_bias = up_down > (T)0.0 ? (T)0.0 : up_down; T down_bias = up_down > (T)0.0 ? -up_down : (T)0.0; + // Infinite granularity mode: when dw_min=0, use dw_min=1 for scale + // computation (unit response) and disable dw_min d-to-d variation. + T effective_dw_min = par.dw_min; T gain_std = par.dw_min_dtod; + if (par.dw_min == (T)0.0) { + effective_dw_min = (T)1.0; + gain_std = (T)0.0; + } // par.w_min = -(T)fabsf(par.w_min); // par.w_max = (T)fabsf(par.w_max); @@ -625,8 +666,8 @@ void PulsedRPUDevice::populate(const PulsedRPUDeviceMetaParameter &p, Real } T r = up_down_std * rng->sampleGauss(); - w_scale_up_[i][j] = (up_bias + gain + r) * par.dw_min; // to reduce mults in updates - w_scale_down_[i][j] = (down_bias + gain - r) * par.dw_min; + w_scale_up_[i][j] = (up_bias + gain + r) * effective_dw_min; // to reduce mults in updates + w_scale_down_[i][j] = (down_bias + gain - r) * effective_dw_min; // enforce consistency if (par.enforce_consistency) { @@ -697,8 +738,8 @@ void PulsedRPUDevice::populate(const PulsedRPUDeviceMetaParameter &p, Real // perfect bias if ((par.perfect_bias) && (j == this->x_size_ - 1)) { - w_scale_up_[i][j] = par.dw_min; - w_scale_down_[i][j] = par.dw_min; + w_scale_up_[i][j] = effective_dw_min; + w_scale_down_[i][j] = effective_dw_min; w_min_bound_[i][j] = (T)100. * par.w_min; // essentially no bound w_max_bound_[i][j] = (T)100. * par.w_max; // essentially no bound } diff --git a/src/rpucuda/rpu_pulsed_device.h b/src/rpucuda/rpu_pulsed_device.h index dfcc2707e..cae5bedd2 100644 --- a/src/rpucuda/rpu_pulsed_device.h +++ b/src/rpucuda/rpu_pulsed_device.h @@ -151,6 +151,16 @@ template class PulsedRPUDeviceBase : public SimpleRPUDevice { virtual void doDenseUpdate(T **weights, int *coincidences, RNG *rng) { RPU_FATAL("Dense update not available for this device!"); }; + virtual void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + RPU_FATAL("Infinite granularity update not available for this device!"); + }; // for Meta-devices [like vector/transfer]: called once before each update starts virtual void initUpdateCycle( T **weights, @@ -275,6 +285,16 @@ template class PulsedRPUDevice : public PulsedRPUDeviceBase { will be re-drawn with noise (even if they were not e.g. clipped) */ + void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) + override; + void decayWeights(T **weights, bool bias_no_decay) override; void decayWeights(T **weights, T alpha, bool bias_no_decay) override; void driftWeights(T **weights, T time_since_last_call, RNG &rng) override; @@ -473,4 +493,19 @@ public: } \ } +// Macro for infinite-granularity (IG) update inner j-loop. +// Computes G = d_val * x_input[j], derives sign and abs_G_lr = lr * |G|. +// BODY should scale device params by abs_G_lr and call update_once with dw_min_std=0. +#define IG_UPDATE_W_LOOP_INNER(BODY) \ + PRAGMA_SIMD \ + for (int j = 0; j < this->x_size_; j++) { \ + T G = d_val * x_input[j * x_inc]; \ + if (G == (T)0.0) { \ + continue; \ + } \ + int sign = (G > (T)0.0) ? 1 : -1; \ + T abs_G_lr = learning_rate * ((G > (T)0.0) ? G : -G); \ + { BODY; } \ + } + } // namespace RPU diff --git a/src/rpucuda/rpu_softbounds_reference_device.cpp b/src/rpucuda/rpu_softbounds_reference_device.cpp index a55deb83f..22637a957 100644 --- a/src/rpucuda/rpu_softbounds_reference_device.cpp +++ b/src/rpucuda/rpu_softbounds_reference_device.cpp @@ -226,6 +226,45 @@ void SoftBoundsReferenceRPUDevice::doDenseUpdate(T **weights, int *coincidenc } } +/* infinite granularity update */ +template +void SoftBoundsReferenceRPUDevice::doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) { + + const auto &par = getPar(); + T write_noise_std = par.write_noise_std * (par.dw_min == (T)0.0 ? (T)1.0 : par.dw_min); + + for (int i = 0; i < this->d_size_; i++) { + T d_val = d_input[i * d_inc]; + if (d_val == (T)0.0) { + continue; + } + + T *scale_down = this->w_scale_down_[i]; + T *scale_up = this->w_scale_up_[i]; + T *ref = w_reference_[i]; + T *w = par.usesPersistentWeight() ? this->w_persistent_[i] : weights[i]; + T *w_apparent = weights[i]; + T *min_bound = this->w_min_bound_[i]; + T *max_bound = this->w_max_bound_[i]; + + IG_UPDATE_W_LOOP_INNER( + T sd = scale_down[j] * abs_G_lr; + T su = scale_up[j] * abs_G_lr; + update_once_mult( + w[j], w_apparent[j], sign, sd, su, ref[j], + min_bound[j], max_bound[j], (T)0.0, write_noise_std, + rng); + ); + } +} + template class SoftBoundsReferenceRPUDevice; #ifdef RPU_USE_DOUBLE template class SoftBoundsReferenceRPUDevice; diff --git a/src/rpucuda/rpu_softbounds_reference_device.h b/src/rpucuda/rpu_softbounds_reference_device.h index f2a5b561a..561738600 100644 --- a/src/rpucuda/rpu_softbounds_reference_device.h +++ b/src/rpucuda/rpu_softbounds_reference_device.h @@ -100,6 +100,15 @@ template class SoftBoundsReferenceRPUDevice : public PulsedRPUDevic T **weights, int i, const int *x_signed_indices, int x_count, int d_sign, RNG *rng) override; void doDenseUpdate(T **weights, int *coincidences, RNG *rng) override; + void doInfiniteGranularityUpdate( + T **weights, + const T *x_input, + int x_inc, + const T *d_input, + int d_inc, + T learning_rate, + RNG *rng) + override; private: T **w_reference_ = nullptr; diff --git a/src/rpucuda/rpu_weight_updater.cpp b/src/rpucuda/rpu_weight_updater.cpp index 747ac3e53..bc80df453 100644 --- a/src/rpucuda/rpu_weight_updater.cpp +++ b/src/rpucuda/rpu_weight_updater.cpp @@ -228,6 +228,18 @@ void PulsedRPUWeightUpdater::updateVectorWithDevice( // check learning rate and update management T weight_granularity = rpu_device->getWeightGranularity(); + // Infinite granularity mode (dw_min=0): exact mean-field update + if (weight_granularity <= (T)0.0) { + rpu_device->initUpdateCycle( + weights, up_, learning_rate, m_batch_info, x_input, x_inc, d_input, d_inc); + T pc_learning_rate = rpu_device->getPulseCountLearningRate(learning_rate, m_batch_info, up_); + T abs_lr = pc_learning_rate < (T)0.0 ? -pc_learning_rate : pc_learning_rate; + rpu_device->doInfiniteGranularityUpdate( + weights, x_input, x_inc, d_input, d_inc, abs_lr, &*rng_); + rpu_device->finishUpdateCycle(weights, up_, learning_rate, m_batch_info); + return; + } + // pulsed device update if (up_.d_sparsity) { up_._d_sparsity = getCurrentDSparsity(); diff --git a/tests/test_infinite_granularity.py b/tests/test_infinite_granularity.py new file mode 100644 index 000000000..ab2c8142e --- /dev/null +++ b/tests/test_infinite_granularity.py @@ -0,0 +1,590 @@ +# -*- coding: utf-8 -*- + +# (C) Copyright 2020, 2021, 2022, 2023, 2024 IBM. All Rights Reserved. +# +# Licensed under the MIT license. See LICENSE file in the project root for details. + +"""Tests for infinite-granularity update mode (dw_min=0).""" + +import time +from unittest import TestCase, skipIf + +import torch +from torch import randn, manual_seed + +from aihwkit.exceptions import ConfigError +from aihwkit.simulator.configs.configs import SingleRPUConfig, UnitCellRPUConfig +from aihwkit.simulator.configs.devices import ( + ConstantStepDevice, + LinearStepDevice, + SoftBoundsDevice, + ExpStepDevice, + PowStepDevice, + PiecewiseStepDevice, + SoftBoundsReferenceDevice, + PowStepReferenceDevice, + TransferCompound, + BufferedTransferCompound, + ChoppedTransferCompound, + DynamicTransferCompound, +) + +SKIP_CUDA = not torch.cuda.is_available() + +DW_MIN_CONVERGENCE_SCAN = [0.001, 0.01, 0.1] + + +IG_DEVICE_CASES = [ + (ConstantStepDevice, {}), + (LinearStepDevice, {"gamma_up": 0.5, "gamma_down": 0.5}), + (SoftBoundsDevice, {}), + (ExpStepDevice, {}), + (PowStepDevice, {"pow_gamma": 0.5}), + ( + PiecewiseStepDevice, + {"piecewise_up": [1.0, 0.5, 1.0], "piecewise_down": [1.0, 0.5, 1.0]}, + ), + (SoftBoundsReferenceDevice, {}), + (PowStepReferenceDevice, {"pow_gamma": 0.5}), +] + +WRITE_NOISE_DEVICE_CASES = [ + (LinearStepDevice, {"gamma_up": 0.5, "gamma_down": 0.5}), + (SoftBoundsDevice, {}), + (ExpStepDevice, {}), + (PowStepDevice, {"pow_gamma": 0.5}), + ( + PiecewiseStepDevice, + {"piecewise_up": [1.0, 0.5, 1.0], "piecewise_down": [1.0, 0.5, 1.0]}, + ), + (SoftBoundsReferenceDevice, {}), +] + + +def _make_tile(device_cls, dw_min, lr=0.01, out_size=5, in_size=4, bias=False, **kwargs): + """Create an AnalogTile with the given device config.""" + from aihwkit.simulator.tiles.analog import AnalogTile + + device = device_cls(dw_min=dw_min, **kwargs) + rpu_config = SingleRPUConfig(device=device) + tile = AnalogTile(out_size, in_size, rpu_config, bias=bias) + tile.tile.set_learning_rate(lr) + return tile + + +def _set_weights_and_update(tile, w_init, x, d): + """Set weights, run one update, return new weights.""" + tile.tile.set_weights(w_init.clone()) + tile.update(x, d) + return tile.tile.get_weights() + + +class InfiniteGranularityBasicTest(TestCase): + """Basic tests that IG mode activates and runs for each device type.""" + + def test_all_ig_capable_devices_run(self): + """All IG-capable concrete device families should run with dw_min=0.""" + manual_seed(1) + for device_cls, kwargs in IG_DEVICE_CASES: + with self.subTest(device=device_cls.__name__): + tile = _make_tile( + device_cls, dw_min=0.0, lr=0.02, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + **kwargs, + ) + w_init = randn(5, 4) * 0.1 + tile.tile.set_weights(w_init.clone()) + tile.update(randn(3, 4), randn(3, 5)) + w_after = tile.tile.get_weights() + self.assertFalse(torch.isnan(w_after).any()) + + def test_constant_step_dw_min_zero(self): + """ConstantStepDevice with dw_min=0 should produce correct update.""" + manual_seed(42) + tile = _make_tile( + ConstantStepDevice, dw_min=0.0, lr=0.1, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + w_init = randn(5, 4) * 0.1 + tile.tile.set_weights(w_init.clone()) + x = randn(3, 4) + d = randn(3, 5) + tile.update(x, d) + w_after = tile.tile.get_weights() + + # Expected: w -= lr * d^T @ x * 1.0 (no variation, beta=0) + G = d.T @ x + w_expected = w_init - 0.1 * G + w_expected = w_expected.clamp(min=-0.6, max=0.6) + torch.testing.assert_close(w_after, w_expected, atol=1e-5, rtol=1e-5) + + def test_constant_step_zero_initial_weights_match_matmul(self): + """Zero-initialized ConstantStepDevice IG update should match matmul.""" + manual_seed(42) + lr = 0.01 + batch_size = 32 + out_size, in_size = 24, 16 + tile = _make_tile( + ConstantStepDevice, dw_min=0.0, lr=lr, + out_size=out_size, in_size=in_size, + w_max=1.0, w_min=-1.0, + w_max_dtod=0.0, w_min_dtod=0.0, + up_down_dtod=0.0, dw_min_dtod=0.0, dw_min_std=0.0, + construction_seed=42, + ) + w_init = torch.zeros(out_size, in_size) + x = randn(batch_size, in_size) * 0.1 + d = randn(batch_size, out_size) * 0.1 + + w_after = _set_weights_and_update(tile, w_init, x, d) + + w_expected = -lr * (d.T @ x) + torch.testing.assert_close(w_after, w_expected, atol=1e-9, rtol=1e-6) + + def test_dw_min_nonzero_unchanged(self): + """dw_min > 0 should still use stochastic update.""" + tile = _make_tile(ConstantStepDevice, dw_min=0.001) + w_init = randn(5, 4) * 0.1 + tile.tile.set_weights(w_init.clone()) + x = randn(3, 4) + d = randn(3, 5) + tile.update(x, d) + w_after = tile.tile.get_weights() + # Should have changed (stochastic) + self.assertFalse(torch.allclose(w_after, w_init, atol=1e-7)) + + def test_linear_step_dw_min_zero(self): + """LinearStepDevice with dw_min=0 should run without errors.""" + tile = _make_tile( + LinearStepDevice, dw_min=0.0, + gamma_up=0.5, gamma_down=0.5, + gamma_up_dtod=0.0, gamma_down_dtod=0.0, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + w_init = randn(5, 4) * 0.1 + tile.tile.set_weights(w_init.clone()) + x = randn(3, 4) + d = randn(3, 5) + tile.update(x, d) + w_after = tile.tile.get_weights() + self.assertFalse(torch.allclose(w_after, w_init, atol=1e-7)) + + def test_soft_bounds_dw_min_zero(self): + """SoftBoundsDevice with dw_min=0 should run.""" + tile = _make_tile(SoftBoundsDevice, dw_min=0.0) + w_init = randn(5, 4) * 0.1 + tile.tile.set_weights(w_init.clone()) + tile.update(randn(3, 4), randn(3, 5)) + + def test_exp_step_dw_min_zero(self): + """ExpStepDevice with dw_min=0 should run.""" + tile = _make_tile(ExpStepDevice, dw_min=0.0) + w_init = randn(5, 4) * 0.1 + tile.tile.set_weights(w_init.clone()) + tile.update(randn(3, 4), randn(3, 5)) + + def test_pow_step_dw_min_zero(self): + """PowStepDevice with dw_min=0 should run.""" + tile = _make_tile(PowStepDevice, dw_min=0.0) + w_init = randn(5, 4) * 0.1 + tile.tile.set_weights(w_init.clone()) + tile.update(randn(3, 4), randn(3, 5)) + + def test_piecewise_step_dw_min_zero(self): + """PiecewiseStepDevice with dw_min=0 should run.""" + tile = _make_tile( + PiecewiseStepDevice, dw_min=0.0, + piecewise_up=[1.0, 0.5, 1.0], + piecewise_down=[1.0, 0.5, 1.0], + ) + w_init = randn(5, 4) * 0.1 + tile.tile.set_weights(w_init.clone()) + tile.update(randn(3, 4), randn(3, 5)) + + +class InfiniteGranularityResponseTest(TestCase): + """Verify response function correctness for specific device types.""" + + def test_linear_step_response_formula(self): + """Verify LinearStepDevice IG update matches the known formula. + + C++ computes slopes as: + slope_up = -ls_decrease_up * scale_up / w_ref_up + slope_down = -ls_decrease_down * scale_down / w_ref_down + With dw_min=0 (effective=1), scale_up=scale_down=1, and + mean_bound_reference=True: w_ref_up=w_max, w_ref_down=w_min. + """ + manual_seed(42) + lr = 0.05 + ls_decrease_up = 0.8 + ls_decrease_down = 0.6 + w_max = 0.6 + w_min = -0.6 + + tile = _make_tile( + LinearStepDevice, dw_min=0.0, lr=lr, + gamma_up=ls_decrease_up, gamma_down=ls_decrease_down, + gamma_up_dtod=0.0, gamma_down_dtod=0.0, + up_down=0.0, up_down_dtod=0.0, + w_max=w_max, w_min=w_min, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + w_init = randn(5, 4) * 0.1 + x = randn(1, 4) # batch=1 for exact formula match + d = randn(1, 5) + w_after = _set_weights_and_update(tile, w_init, x, d) + + # Manually compute expected result matching C++ slope formula + G = d.T @ x + scale_up = 1.0 # effective_dw_min=1, gain=1, up_bias=0 + scale_down = 1.0 + slope_up = -ls_decrease_up * scale_up / w_max + slope_down = -ls_decrease_down * scale_down / w_min + + resp_up = slope_up * w_init + scale_up + resp_down = slope_down * w_init + scale_down + resp = torch.where(G > 0, resp_down, resp_up) + + w_expected = w_init - lr * G * resp + w_expected = w_expected.clamp(min=w_min, max=w_max) + + torch.testing.assert_close(w_after, w_expected, atol=1e-5, rtol=1e-5) + + +class InfiniteGranularityConvergenceTest(TestCase): + """Convergence: avg stochastic update stays within the dw_min scale.""" + + def _convergence_test(self, device_cls, dw_min, n_repeat=1000, **kwargs): + manual_seed(0) + lr = 0.1 + out_size, in_size = 3, 3 + common_kwargs = dict( + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + dw_min_dtod=0.0, dw_min_std=0.0, + **kwargs, + ) + + w_init = randn(out_size, in_size) * 0.1 + x = randn(1, in_size) * 0.3 + d = randn(1, out_size) * 0.3 + + tile_ig = _make_tile( + device_cls, dw_min=0.0, lr=lr, + out_size=out_size, in_size=in_size, **common_kwargs + ) + w_ig = _set_weights_and_update(tile_ig, w_init, x, d) + + w_sum = torch.zeros_like(w_init) + for _ in range(n_repeat): + tile_s = _make_tile( + device_cls, dw_min=dw_min, lr=lr, + out_size=out_size, in_size=in_size, **common_kwargs + ) + w_s = _set_weights_and_update(tile_s, w_init, x, d) + w_sum += w_s + w_avg = w_sum / n_repeat + + err = (w_avg - w_ig).abs().mean().item() + self.assertLess( + err, dw_min, + f"Convergence failed for {device_cls.__name__}: " + f"dw_min={dw_min:.6f}, err={err:.6f}" + ) + + def test_convergence_constant_step(self): + for dw_min in DW_MIN_CONVERGENCE_SCAN: + with self.subTest(dw_min=dw_min): + self._convergence_test(ConstantStepDevice, dw_min=dw_min) + + def test_convergence_linear_step(self): + for dw_min in DW_MIN_CONVERGENCE_SCAN: + with self.subTest(dw_min=dw_min): + self._convergence_test( + LinearStepDevice, + dw_min=dw_min, + gamma_up=0.5, gamma_down=0.5, + gamma_up_dtod=0.0, gamma_down_dtod=0.0, + ) + + def test_convergence_soft_bounds(self): + for dw_min in DW_MIN_CONVERGENCE_SCAN: + with self.subTest(dw_min=dw_min): + self._convergence_test(SoftBoundsDevice, dw_min=dw_min) + + +class InfiniteGranularityWriteNoiseTest(TestCase): + """Smoke tests for IG persistent/apparent write-noise path.""" + + def _write_noise_test(self, device_cls, cuda=False, **kwargs): + manual_seed(7) + tile = _make_tile( + device_cls, dw_min=0.0, lr=0.01, + out_size=5, in_size=4, + write_noise_std=0.01, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + **kwargs, + ) + if cuda: + tile = tile.cuda() + + w_init = randn(5, 4) * 0.05 + x = randn(2, 4) * 0.1 + d = randn(2, 5) * 0.1 + tile.tile.set_weights(w_init.clone()) + tile.update(x.cuda() if cuda else x, d.cuda() if cuda else d) + + w_after = tile.tile.get_weights() + hidden = tile.get_hidden_parameters() + self.assertIn("persistent_weights", hidden) + persistent = hidden["persistent_weights"] + + self.assertFalse(torch.isnan(w_after).any()) + self.assertGreater((w_after.cpu() - persistent.cpu()).abs().max().item(), 1e-6) + self.assertLessEqual(persistent.max().item(), 0.6 + 1e-6) + self.assertGreaterEqual(persistent.min().item(), -0.6 - 1e-6) + + def test_cpu_write_noise_devices(self): + for device_cls, kwargs in WRITE_NOISE_DEVICE_CASES: + with self.subTest(device=device_cls.__name__): + self._write_noise_test(device_cls, **kwargs) + + @skipIf(SKIP_CUDA, "CUDA not available") + def test_gpu_write_noise_devices(self): + for device_cls, kwargs in WRITE_NOISE_DEVICE_CASES: + with self.subTest(device=device_cls.__name__): + self._write_noise_test(device_cls, cuda=True, **kwargs) + + +@skipIf(SKIP_CUDA, "CUDA not available") +class InfiniteGranularityCPUGPUConsistencyTest(TestCase): + """Verify CPU and GPU produce identical results for IG mode.""" + + def _cpu_gpu_test(self, device_cls, **kwargs): + """Test CPU/GPU consistency with batch_size=1.""" + manual_seed(42) + lr = 0.05 + out_size, in_size = 8, 6 + + tile_cpu = _make_tile(device_cls, dw_min=0.0, lr=lr, + out_size=out_size, in_size=in_size, **kwargs) + tile_gpu = _make_tile(device_cls, dw_min=0.0, lr=lr, + out_size=out_size, in_size=in_size, **kwargs) + tile_gpu = tile_gpu.cuda() + + w_init = randn(out_size, in_size) * 0.1 + x = randn(1, in_size) + d = randn(1, out_size) + + tile_cpu.tile.set_weights(w_init.clone()) + tile_gpu.tile.set_weights(w_init.clone()) + + tile_cpu.update(x, d) + tile_gpu.update(x.cuda(), d.cuda()) + + w_cpu = tile_cpu.tile.get_weights() + w_gpu = tile_gpu.tile.get_weights() + + torch.testing.assert_close( + w_cpu, w_gpu, atol=1e-5, rtol=1e-5, + msg=f"CPU/GPU mismatch for {device_cls.__name__}" + ) + + def test_cpu_gpu_constant_step(self): + self._cpu_gpu_test( + ConstantStepDevice, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + def test_cpu_gpu_linear_step(self): + self._cpu_gpu_test( + LinearStepDevice, + gamma_up=0.5, gamma_down=0.5, + gamma_up_dtod=0.0, gamma_down_dtod=0.0, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + def test_cpu_gpu_soft_bounds(self): + self._cpu_gpu_test( + SoftBoundsDevice, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + def test_cpu_gpu_exp_step(self): + self._cpu_gpu_test( + ExpStepDevice, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + def test_cpu_gpu_pow_step(self): + self._cpu_gpu_test( + PowStepDevice, + pow_gamma=0.5, pow_gamma_dtod=0.0, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + def test_cpu_gpu_piecewise_step(self): + self._cpu_gpu_test( + PiecewiseStepDevice, + piecewise_up=[1.0, 0.8, 0.6, 0.4, 0.2], + piecewise_down=[0.2, 0.4, 0.6, 0.8, 1.0], + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + def test_cpu_gpu_softbounds_reference(self): + self._cpu_gpu_test( + SoftBoundsReferenceDevice, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + def test_cpu_gpu_powstep_reference(self): + self._cpu_gpu_test( + PowStepReferenceDevice, + pow_gamma=0.5, pow_gamma_dtod=0.0, + up_down=0.0, up_down_dtod=0.0, + w_max_dtod=0.0, w_min_dtod=0.0, + ) + + +@skipIf(SKIP_CUDA, "CUDA not available") +class InfiniteGranularityPerformanceTest(TestCase): + """Performance and memory benchmarks for IG mode.""" + + def _benchmark(self, device_cls, dw_min, out_size, in_size, batch_size, n_iter=50, **kwargs): + """Run n_iter updates and return (total_time, peak_memory_bytes).""" + rpu_config = SingleRPUConfig(device=device_cls(dw_min=dw_min, **kwargs)) + from aihwkit.simulator.tiles.analog import AnalogTile + tile = AnalogTile(out_size, in_size, rpu_config, bias=False) + tile.tile.set_learning_rate(0.01) + tile = tile.cuda() + + w_init = randn(out_size, in_size) * 0.1 + tile.tile.set_weights(w_init) + + x = randn(batch_size, in_size).cuda() + d = randn(batch_size, out_size).cuda() + + # Warmup + for _ in range(5): + tile.tile.set_weights(w_init) + tile.update(x, d) + torch.cuda.synchronize() + + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + start = time.perf_counter() + for _ in range(n_iter): + tile.tile.set_weights(w_init) + tile.update(x, d) + torch.cuda.synchronize() + elapsed = time.perf_counter() - start + peak_mem = torch.cuda.max_memory_allocated() + + return elapsed, peak_mem + + def test_performance_small_batch(self): + """IG should be comparable or faster than stochastic for small batch.""" + for batch_size in [1, 8]: + t_ig, _ = self._benchmark(ConstantStepDevice, dw_min=0.0, out_size=64, in_size=64, + batch_size=batch_size) + t_st, _ = self._benchmark(ConstantStepDevice, dw_min=0.001, out_size=64, in_size=64, + batch_size=batch_size) + print(f" batch={batch_size}: IG={t_ig:.4f}s, stochastic={t_st:.4f}s, " + f"ratio={t_ig / t_st:.2f}") + + def test_performance_large_batch(self): + """IG with large batch should not be excessively slow.""" + for batch_size in [256, 1024, 4096]: + t_ig, _ = self._benchmark(ConstantStepDevice, dw_min=0.0, out_size=64, in_size=64, + batch_size=batch_size) + t_st, _ = self._benchmark(ConstantStepDevice, dw_min=0.001, out_size=64, in_size=64, + batch_size=batch_size) + print(f" batch={batch_size}: IG={t_ig:.4f}s, stochastic={t_st:.4f}s, " + f"ratio={t_ig / t_st:.2f}") + + def test_memory_large_batch(self): + """IG should not use significantly more memory than stochastic.""" + for batch_size in [1024, 4096, 16384]: + _, mem_ig = self._benchmark(ConstantStepDevice, dw_min=0.0, out_size=128, in_size=128, + batch_size=batch_size, n_iter=10) + _, mem_st = self._benchmark(ConstantStepDevice, dw_min=0.001, out_size=128, in_size=128, + batch_size=batch_size, n_iter=10) + diff_mb = (mem_ig - mem_st) / (1024 * 1024) + print(f" batch={batch_size}: IG_mem={mem_ig / (1024 * 1024):.1f}MB, " + f"stochastic_mem={mem_st / (1024 * 1024):.1f}MB, diff={diff_mb:.1f}MB") + # IG should not use significantly more memory than stochastic + # The IG GEMM allocates input buffers + grad matrix temp (~same as stochastic) + # Allow up to 20MB extra + self.assertLess( + diff_mb, 20.0, + f"IG uses too much extra memory at batch={batch_size}: " + f"diff={diff_mb:.1f}MB (IG={mem_ig}, stochastic={mem_st})" + ) + + +class InfiniteGranularityTransferGuardTest(TestCase): + """IG mode (dw_min=0) must be rejected inside transfer compounds. + + Transfer compounds scale the transfer learning rate and buffer + granularity by the sub-device weight granularity (dw_min); a zero + dw_min would silently disable learning, so it must raise a ConfigError. + """ + + TRANSFER_COMPOUNDS = [ + TransferCompound, + BufferedTransferCompound, + ChoppedTransferCompound, + DynamicTransferCompound, + ] + + def test_dw_min_zero_rejected_on_construction(self): + """Constructing a transfer compound with a dw_min=0 sub-device raises.""" + for compound_cls in self.TRANSFER_COMPOUNDS: + with self.subTest(compound=compound_cls.__name__): + with self.assertRaises(ConfigError): + compound_cls( + unit_cell_devices=[ + ConstantStepDevice(dw_min=0.0), + ConstantStepDevice(dw_min=0.001), + ] + ) + + def test_dw_min_zero_on_slow_device_rejected(self): + """A dw_min=0 on any sub-device (not only the fast one) raises.""" + with self.assertRaises(ConfigError): + ChoppedTransferCompound( + unit_cell_devices=[ + ConstantStepDevice(dw_min=0.001), + ConstantStepDevice(dw_min=0.0), + ] + ) + + def test_nonzero_dw_min_transfer_compound_builds(self): + """A transfer compound with non-zero dw_min builds a working tile.""" + rpu_config = UnitCellRPUConfig( + device=ChoppedTransferCompound( + unit_cell_devices=[ + ConstantStepDevice(dw_min=0.001), + ConstantStepDevice(dw_min=0.001), + ] + ) + ) + tile = rpu_config.tile_class(5, 4, rpu_config, bias=False) + tile.update(randn(3, 4), randn(3, 5)) + + def test_single_device_dw_min_zero_still_allowed(self): + """IG mode remains valid for a plain (non-compound) pulsed device.""" + tile = _make_tile(ConstantStepDevice, dw_min=0.0) + tile.update(randn(3, 4), randn(3, 5))