refactor(mge/examples): add example to use modified sublinear API

GitOrigin-RevId: c6f6c65d0bf2011bb8590aec3285bbac99fc9c34

python_module/megengine/jit/__init__.py

@@ -18,7 +18,7 @@ import megengine._internal as mgb
 from megengine._internal.plugin import CompGraphProfiler
 
 from ..core import Tensor, graph, tensor
-from .sublinear_memory_config import SublinearMemConfig
+from .sublinear_memory_config import SublinearMemoryConfig
 
 
 def sideeffect(f):
@@ -82,8 +82,8 @@ class trace:
     :param symbolic: Whether to use symbolic tensor. Default: False
     :param opt_level: Optimization level for compiling trace.
     :param log_level: Log level.
-    :param enable_sublinear: Enable sublinear memory optimization. Default: False
-    :param sublinear_mem_config: Configuration for sublinear memory optimization.
+    :param sublinear_memory_config: Configuration for sublinear memory optimization.
+        If not None, it enables sublinear memory optimization with given setting.
     :param profiling: Whether to profile compiled trace. Default: False
     """
 
@@ -106,16 +106,14 @@ class trace:
         symbolic: bool = False,
         opt_level: int = None,
         log_level: int = None,
-        enable_sublinear: bool = False,
-        sublinear_mem_config: SublinearMemConfig = None,
+        sublinear_memory_config: SublinearMemoryConfig = None,
         profiling: bool = False
     ):
         self.__wrapped__ = func
         self._symbolic = symbolic
         self._graph_opt_level = opt_level
         self._log_level = log_level
-        self._enable_sublinear = enable_sublinear
-        self._sublinear_mem_config = sublinear_mem_config
+        self._sublinear_memory_config = sublinear_memory_config
         self._status = self._UNSTARTED
         self._args = None
         self._kwargs = None
@@ -287,35 +285,34 @@ class trace:
 
     def _apply_graph_options(self, cg):
         # graph opt level
-        if not (self._graph_opt_level is None):
+        if self._graph_opt_level is not None:
             cg.set_option("graph_opt_level", self._graph_opt_level)
         # log level
-        if not (self._log_level is None):
+        if self._log_level is not None:
             cg.set_option("log_level", self._log_level)
         # sublinear
-        if self._enable_sublinear:
+        if self._sublinear_memory_config is not None:
             cg.set_option("enable_sublinear_memory_opt", True)
-            if not (self._sublinear_mem_config is None):
-                cg.set_option(
-                    "sublinear_mem_cofig.lb_memory",
-                    self._sublinear_mem_config.lb_memory,
-                )
-                cg.set_option(
-                    "sublinear_mem_cofig.genetic_nr_iter",
-                    self._sublinear_mem_config.genetic_nr_iter,
-                )
-                cg.set_option(
-                    "sublinear_mem_cofig.genetic_pool_size",
-                    self._sublinear_mem_config.genetic_pool_size,
-                )
-                cg.set_option(
-                    "sublinear_mem_cofig.thresh_nr_try",
-                    self._sublinear_mem_config.thresh_nr_try,
-                )
-                cg.set_option(
-                    "sublinear_mem_cofig.num_worker",
-                    self._sublinear_mem_config.num_worker,
-                )
+            cg.set_option(
+                "sublinear_mem_cofig.lb_memory",
+                self._sublinear_memory_config.lb_memory,
+            )
+            cg.set_option(
+                "sublinear_mem_cofig.genetic_nr_iter",
+                self._sublinear_memory_config.genetic_nr_iter,
+            )
+            cg.set_option(
+                "sublinear_mem_cofig.genetic_pool_size",
+                self._sublinear_memory_config.genetic_pool_size,
+            )
+            cg.set_option(
+                "sublinear_mem_cofig.thresh_nr_try",
+                self._sublinear_memory_config.thresh_nr_try,
+            )
+            cg.set_option(
+                "sublinear_mem_cofig.num_worker",
+                self._sublinear_memory_config.num_worker,
+            )
         # profile
         if self._profiling:
             self._profiler = CompGraphProfiler(cg)

python_module/megengine/jit/sublinear_memory_config.py

@@ -10,7 +10,7 @@
 from ..core.device import get_device_count
 
 
-class SublinearMemConfig:
+class SublinearMemoryConfig:
     r"""
     Configuration for sublinear memory optimization.
 
@@ -28,7 +28,11 @@ class SublinearMemConfig:
         It can also be set through the environmental variable 'MGB_SUBLINEAR_MEMORY_LOWER_BOUND_MB'.
     :param num_worker: number of thread workers to search the optimum checkpoints
         in sublinear memory optimization. Default: half of cpu number in the system.
+        Note: the value must be greater or equal to one.
         It can also be set through the environmental variable 'MGB_SUBLINEAR_MEMORY_WORKERS'.
+
+    Note that the environmental variable MGB_COMP_GRAPH_OPT must be set to 'enable_sublinear_memory_opt=1'
+    in order for the above environmental variable to be effective.
     """
 
     def __init__(
@@ -37,10 +41,16 @@ class SublinearMemConfig:
         genetic_nr_iter: int = 0,
         genetic_pool_size: int = 20,
         lb_memory: int = 0,
-        num_worker: int = get_device_count("cpu") / 2,
+        num_worker: int = max(1, get_device_count("cpu") // 2),
     ):
+        assert thresh_nr_try >= 0, "thresh_nr_try must be greater or equal to zero"
         self.thresh_nr_try = thresh_nr_try
+        assert genetic_nr_iter >= 0, "genetic_nr_iter must be greater or equal to zero"
         self.genetic_nr_iter = genetic_nr_iter
+        assert (
+            genetic_pool_size >= 0
+        ), "genetic_pool_size must be greater or equal to zero"
         self.genetic_pool_size = genetic_pool_size
         self.lb_memory = lb_memory
+        assert num_worker > 0, "num_worker must be greater or equal to one"
         self.num_worker = num_worker

python_module/test/integration/test_correctness.py

@@ -17,7 +17,7 @@ import megengine as mge
 import megengine.functional as F
 from megengine import jit, tensor
 from megengine.functional.debug_param import set_conv_execution_strategy
-from megengine.jit import SublinearMemConfig
+from megengine.jit import SublinearMemoryConfig
 from megengine.module import AvgPool2d, BatchNorm2d, Conv2d, Linear, Module
 from megengine.optimizer import SGD
 from megengine.test import assertTensorClose
@@ -132,12 +132,7 @@ def update_model(model_path):
 
 
 def run_test(
-    model_path,
-    use_jit,
-    use_symbolic,
-    enable_sublinear=False,
-    sublinear_mem_config=None,
-    max_err=None,
+    model_path, use_jit, use_symbolic, sublinear_memory_config=None, max_err=None,
 ):
 
     """
@@ -168,8 +163,7 @@ def run_test(
         train_func = jit.trace(
             train_func,
             symbolic=use_symbolic,
-            enable_sublinear=enable_sublinear,
-            sublinear_mem_config=sublinear_mem_config,
+            sublinear_memory_config=sublinear_memory_config,
         )
 
     opt.zero_grad()
@@ -199,12 +193,7 @@ def test_correctness():
     run_test(model_path, True, True)
 
     # sublinear
-    config = SublinearMemConfig(genetic_nr_iter=10)
+    config = SublinearMemoryConfig(genetic_nr_iter=10)
     run_test(
-        model_path,
-        True,
-        True,
-        enable_sublinear=True,
-        sublinear_mem_config=config,
-        max_err=1e-5,
+        model_path, True, True, sublinear_memory_config=config, max_err=1e-5,
     )

python_module/test/unit/jit/test_jit.py

@@ -18,7 +18,7 @@ import megengine._internal as mgb
 import megengine.module as M
 from megengine import jit, tensor
 from megengine.core.tensor import Tensor
-from megengine.jit import SublinearMemConfig
+from megengine.jit import SublinearMemoryConfig
 from megengine.test import assertTensorClose
 
 
@@ -190,9 +190,9 @@ def test_dump_bn_fused():
 
 # Simply verify the options passed down
 def test_sublinear():
-    config = SublinearMemConfig(genetic_nr_iter=10)
+    config = SublinearMemoryConfig(genetic_nr_iter=10)
 
-    @jit.trace(symbolic=True, enable_sublinear=True, sublinear_mem_config=config)
+    @jit.trace(symbolic=True, sublinear_memory_config=config)
     def f(x):
         return x + x