add offload mode

ckpt/songgeneration_base/config.yaml

@@ -106,3 +106,36 @@ conditioners:
     QwTextTokenizer:
       token_path: third_party/Qwen2-7B
       max_len: 50
+
+offload:
+  audiolm:
+    offload_module: self
+    cpu_mem_gb: 0
+    pre_copy_step: 1
+    clean_cache_after_forward: false
+    dtype: torch.float16
+    offload_layer_dict:
+      transformer: 4
+      transformer2: 4
+    ignore_layer_list: []
+    clean_cache_wrapper:
+      module: self
+      method_name: _sample_next_token
+      diff_mem_gb_thre: 2
+    debug: false
+
+  wav_tokenizer_diffusion:
+    offload_module: self.model.model
+    pre_copy_step: 1
+    clean_cache_after_forward: false
+    cpu_mem_gb: -1
+    dtype: null
+    offload_layer_dict:
+      cfm_wrapper: 5
+      hubert: 4
+    ignore_layer_list: []
+    clean_cache_wrapper:
+      module: self.model.model.cfm_wrapper.estimator
+      method_name: forward
+      diff_mem_gb_thre: 1
+    debug: false