Fix chopper bug (IBM#732)

* fix: debug mode, chopper trans bug, memory issue, maximizer trans

* fix memeory error with valgrind

* changelog

* fix UMH test

* CUDA arch

Author: maljoras-sony

CHANGELOG.md

@@ -27,6 +27,7 @@ The format is based on [Keep a Changelog], and this project adheres to
 
 * Fix Hardware-Aware training tutorial notebooks (\#700)
 * Fix Post-Training Input Range Calibration notebook (\#716)
+* Fix memory issues and bugs in analog training for CUDA (\#732)
 
 ## Changed
 

CMakeLists.txt

@@ -23,7 +23,7 @@ option(RPU_USE_TORCH_BUFFERS "Use torch buffers for RPUCuda" ON)
 
 
 set(RPU_BLAS "OpenBLAS" CACHE STRING "BLAS backend of choice (OpenBLAS, MKL)")
-set(RPU_CUDA_ARCHITECTURES "70;75;80" CACHE STRING "Target CUDA architectures")
+set(RPU_CUDA_ARCHITECTURES "75;80;89" CACHE STRING "Target CUDA architectures")
 
 # Internal variables.
 set(CUDA_TARGET_PROPERTIES POSITION_INDEPENDENT_CODE ON
@@ -110,12 +110,13 @@ if(USE_CUDA)
   if (RPU_USE_TORCH_BUFFERS)
     if (BUILD_TEST)
       # we could just link torch to the tests in principle
-      message(FATAL_ERROR "Cannot use torch buffers when BUILD_TEST=ON. Set RPU_USE_TORCH_BUFFERS=OFF")
+      message(STATUS "Cannot use torch buffers when BUILD_TEST=ON. Set RPU_USE_TORCH_BUFFERS=OFF")
+      set(RPU_USE_TORCH_BUFFERS OFF)
+    else (BUILD_TEST)
+      add_compile_definitions(RPU_TORCH_CUDA_BUFFERS)
+      set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
+      set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcudafe --diag_suppress=186")
     endif(BUILD_TEST)
-
-    add_compile_definitions(RPU_TORCH_CUDA_BUFFERS)
-    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
-    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcudafe --diag_suppress=186")
   endif(RPU_USE_TORCH_BUFFERS)
 
   if(${CUDAToolkit_VERSION_MAJOR} LESS 11)
@@ -140,7 +141,7 @@ if (BUILD_EXTENSION)
 
   target_link_libraries(AIHWKIT_EXTENSION_OPS torch_python c10 torch_cpu)
   target_include_directories(AIHWKIT_EXTENSION_OPS PRIVATE src/aihwkit/extension/extension_src)
-  
+
   if(WIN32)
     target_link_libraries(AIHWKIT_EXTENSION_OPS c10.lib torch_cpu.lib)
   endif()
@@ -149,7 +150,7 @@ if (BUILD_EXTENSION)
     add_library(AIHWKIT_EXTENSION_OPS_GPU ${AIHWKIT_EXTENSION_OPS_GPU_SRCS})
     target_link_libraries(AIHWKIT_EXTENSION_OPS_GPU AIHWKIT_EXTENSION_OPS c10_cuda torch_cuda cudart)
     target_include_directories(AIHWKIT_EXTENSION_OPS_GPU PRIVATE src/aihwkit/extension/extension_src)
-    
+
     set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
     set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcudafe --diag_suppress=186")
     set_target_properties(AIHWKIT_EXTENSION_OPS_GPU PROPERTIES ${CUDA_TARGET_PROPERTIES})
@@ -187,20 +188,34 @@ endif(BUILD_EXTENSION)
 
 # Add tests.
 if(BUILD_TEST)
+
   enable_testing()
 
   foreach(test_src ${RPU_CPU_TEST_SRCS} ${RPU_GPU_TEST_SRCS})
     get_filename_component(test_name ${test_src} NAME_WE)
     add_executable(${test_name} ${test_src})
     target_link_libraries(${test_name} gtest gmock)
+    target_link_libraries(${test_name} torch_python c10 torch_cpu)
+    set_target_properties(${test_name} PROPERTIES CXX_STANDARD 17
+      POSITION_INDEPENDENT_CODE ON)
+
+    if(WIN32)
+      target_link_libraries(${test_name} c10.lib torch_cpu.lib)
+    endif()
 
     # Link to main library.
-    if(${test_src} IN_LIST RPU_CPU_TEST_SRCS)
-      target_link_libraries(${test_name} RPU_CPU ${RPU_DEPENDENCY_LIBS})
-    else()
+    target_link_libraries(${test_name} RPU_CPU ${RPU_DEPENDENCY_LIBS})
+
+    if(${test_src} IN_LIST RPU_GPU_TEST_SRCS)
+      target_link_libraries(${test_name} torch_cuda c10_cuda cudart)
       target_link_libraries(${test_name} RPU_GPU RPU_CPU cublas curand ${RPU_DEPENDENCY_LIBS})
       set_target_properties(${test_name} PROPERTIES ${CUDA_TARGET_PROPERTIES})
       set_property(TARGET ${test_name} PROPERTY CUDA_ARCHITECTURES ${RPU_CUDA_ARCHITECTURES})
+
+      if(WIN32)
+	target_link_libraries(${test_name} c10_cuda.lib torch_cuda.lib)
+      endif(WIN32)
+
     endif()
 
     add_test(NAME ${test_name} COMMAND $<TARGET_FILE:${test_name}>)

cmake/dependencies_test.cmake

@@ -11,6 +11,7 @@ if(BUILD_TEST)
     URL_HASH          MD5=52943a59cefce0ae0491d4d2412c120b
     CMAKE_ARGS        "-DCMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI\=0"
     INSTALL_COMMAND   ""
+    DOWNLOAD_EXTRACT_TIMESTAMP true
   )
 
   ExternalProject_Get_Property(GTest source_dir)
@@ -20,4 +21,5 @@ if(BUILD_TEST)
 
   include_directories(SYSTEM ${GTest_INCLUDE_DIR})
   link_directories(SYSTEM ${GTest_LIBRARY_DIR})
+ 
 endif()

examples/04_lenet5_training.py

@@ -29,14 +29,13 @@
 from aihwkit.nn import AnalogConv2d, AnalogLinear, AnalogSequential
 from aihwkit.optim import AnalogSGD
 from aihwkit.simulator.configs import (
-    SingleRPUConfig,
     FloatingPointRPUConfig,
-    ConstantStepDevice,
+    SoftBoundsReferenceDevice,
     FloatingPointDevice,
+    build_config,
 )
 from aihwkit.simulator.rpu_base import cuda
 
-
 # Check device
 USE_CUDA = 0
 if cuda.is_compiled():
@@ -59,13 +58,17 @@
 # Select the device model to use in the training.
 # * If `SingleRPUConfig(device=ConstantStepDevice())` then analog tiles with
 #   constant step devices will be used,
+# * One can use `build_config` to build a config for different
+#   specialized analog gradient algorithms
 # * If `FloatingPointRPUConfig(device=FloatingPointDevice())` then standard
 #   floating point devices will be used
-USE_ANALOG_TRAINING = False
+USE_ANALOG_TRAINING = True
 if USE_ANALOG_TRAINING:
-    RPU_CONFIG = SingleRPUConfig(device=ConstantStepDevice())
+    algo = "agad"  # or e.g. ttv2
+    RPU_CONFIG = build_config(algo, device=SoftBoundsReferenceDevice(dw_min=0.05))
 else:
     RPU_CONFIG = FloatingPointRPUConfig(device=FloatingPointDevice())
+print(RPU_CONFIG)
 
 
 def load_images():

examples/23_using_analog_tile_as_matrix.py

@@ -4,8 +4,7 @@
 #
 # Licensed under the MIT license. See LICENSE file in the project root for details.
 
-"""aihwkit example 22: Simple example of how to use an analog tile as a matrix
-"""
+"""aihwkit example 22: Simple example of how to use an analog tile as a matrix"""
 # pylint: disable=invalid-name
 # pylint: disable=too-many-locals
 

examples/26_correlation_detection.py

@@ -4,8 +4,7 @@
 #
 # Licensed under the MIT license. See LICENSE file in the project root for details.
 
-"""aihwkit example 25: Simple correlation detection with analog optimizers.
-"""
+"""aihwkit example 25: Simple correlation detection with analog optimizers."""
 # pylint: disable=invalid-name, too-many-locals, too-many-statements
 
 from typing import Union, Tuple, Optional, List, Dict

examples/35_half_precision_training.py

@@ -70,9 +70,7 @@ def forward(self, x):
 
     pbar = tqdm.tqdm(enumerate(train_loader))
     for batch_idx, (data, target) in pbar:
-        data, target = data.to(device=device, dtype=torch.bfloat16), target.to(
-            device=device
-        )
+        data, target = data.to(device=device, dtype=torch.bfloat16), target.to(device=device)
         optimizer.zero_grad()
         output = model(data)
         loss = F.nll_loss(output.float(), target)

src/aihwkit/extension/__init__.py

@@ -6,7 +6,7 @@
 
 # pylint: disable=import-error, no-name-in-module, invalid-name
 
-"""AIHWKIT extension """
+"""AIHWKIT extension"""
 
 from importlib.util import find_spec
 

src/aihwkit/linalg/matrix.py

@@ -4,8 +4,7 @@
 #
 # Licensed under the MIT license. See LICENSE file in the project root for details.
 
-""" Defines an analog matrix
-"""
+"""Defines an analog matrix"""
 
 from typing import Any, Union, Tuple
 from scipy.sparse.linalg import LinearOperator

src/aihwkit/nn/low_precision_conversion.py

@@ -4,7 +4,7 @@
 #
 # Licensed under the MIT license. See LICENSE file in the project root for details.
 
-""" Functions to convert a given model to a quantized counterpart """
+"""Functions to convert a given model to a quantized counterpart"""
 
 from copy import deepcopy