Repository: voicepaw/so-vits-svc-fork

Issue:
CUDA Out of Memory during training
Hi, I've trained a few models without issue but after updating to 3.0 when trying to train a new model I run out of memory. I'm using the same dataset as before and there's no issues until the SVC Train command. Any idea what might be causing it? I can't make head or tail of the error message.

```
Traceback (most recent call last):
  File "C:\Program Files\Python310\lib\runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "C:\Program Files\Python310\lib\runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\Scripts\svc.exe\__main__.py", line 7, in <module>
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\click\core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\click\core.py", line 1055, in main
    rv = self.invoke(ctx)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\click\core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\click\core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\click\core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\so_vits_svc_fork\__main__.py", line 128, in train
    train(
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\so_vits_svc_fork\train.py", line 53, in train
    mp.spawn(
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\multiprocessing\spawn.py", line 239, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\multiprocessing\spawn.py", line 197, in start_processes
    while not context.join():
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\multiprocessing\spawn.py", line 160, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException:

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\multiprocessing\spawn.py", line 69, in _wrap
    fn(i, *args)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\so_vits_svc_fork\train.py", line 167, in _run
    _train_and_evaluate(
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\so_vits_svc_fork\train.py", line 266, in _train_and_evaluate
    y_d_hat_r, y_d_hat_g, _, _ = net_d(y, y_hat.detach())
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\nn\parallel\distributed.py", line 1156, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\nn\parallel\distributed.py", line 1110, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])  # type: ignore[index]
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\so_vits_svc_fork\modules\descriminators.py", line 137, in forward
    y_d_r, fmap_r = d(y)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\so_vits_svc_fork\modules\descriminators.py", line 81, in forward
    x = F.leaky_relu(x, modules.LRELU_SLOPE)
  File "C:\Users\Star Guard 719\Documents\sovits new\venv\lib\site-packages\torch\nn\functional.py", line 1632, in leaky_relu
    result = torch._C._nn.leaky_relu(input, negative_slope)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 12.00 MiB (GPU 0; 10.00 GiB total capacity; 9.13 GiB already allocated; 0 bytes free; 9.28 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```



Discussion:
Please reduce batch_size

Resolving patch:
diff --git a/src/so_vits_svc_fork/preprocessing/preprocess_hubert_f0.py b/src/so_vits_svc_fork/preprocessing/preprocess_hubert_f0.py
--- a/src/so_vits_svc_fork/preprocessing/preprocess_hubert_f0.py
+++ b/src/so_vits_svc_fork/preprocessing/preprocess_hubert_f0.py
@@ -10,7 +10,7 @@
 import torch
 import torchaudio
 from fairseq.models.hubert import HubertModel
-from joblib import Parallel, delayed
+from joblib import Parallel, cpu_count, delayed
 from tqdm import tqdm
 
 import so_vits_svc_fork.f0
@@ -22,8 +22,8 @@
 from .preprocess_utils import check_hubert_min_duration
 
 LOG = getLogger(__name__)
-HUBERT_MEMORY = 1600
-HUBERT_MEMORY_CREPE = 2600
+HUBERT_MEMORY = 2900
+HUBERT_MEMORY_CREPE = 3900
 
 
 def _process_one(
@@ -124,11 +124,17 @@ def preprocess_hubert_f0(
     utils.ensure_pretrained_model(".", "contentvec")
     hps = utils.get_hparams(config_path)
     if n_jobs is None:
-        memory = get_total_gpu_memory("free")
-        n_jobs = (
-            memory // (HUBERT_MEMORY_CREPE if f0_method == "crepe" else HUBERT_MEMORY)
-            if memory is not None
-            else 1
+        # add cpu_count() to avoid SIGKILL
+        memory = get_total_gpu_memory("total")
+        n_jobs = min(
+            max(
+                memory
+                // (HUBERT_MEMORY_CREPE if f0_method == "crepe" else HUBERT_MEMORY)
+                if memory is not None
+                else 1,
+                1,
+            ),
+            cpu_count(),
         )
         LOG.info(f"n_jobs automatically set to {n_jobs}, memory: {memory} MiB")
 
diff --git a/src/so_vits_svc_fork/utils.py b/src/so_vits_svc_fork/utils.py
--- a/src/so_vits_svc_fork/utils.py
+++ b/src/so_vits_svc_fork/utils.py
@@ -163,6 +163,32 @@ def get_content(
     return c
 
 
+def _substitute_if_same_shape(to_: dict[str, Any], from_: dict[str, Any]) -> None:
+    for k, v in from_.items():
+        if k not in to_:
+            warnings.warn(f"Key {k} not found in model state dict")
+        elif hasattr(v, "shape"):
+            if not hasattr(to_[k], "shape"):
+                raise ValueError(f"Key {k} is not a tensor")
+            if to_[k].shape == v.shape:
+                to_[k] = v
+            else:
+                warnings.warn(
+                    f"Shape mismatch for key {k}, {to_[k].shape} != {v.shape}"
+                )
+        elif isinstance(v, dict):
+            assert isinstance(to_[k], dict)
+            _substitute_if_same_shape(to_[k], v)
+        else:
+            to_[k] = v
+
+
+def safe_load(model: torch.nn.Module, state_dict: dict[str, Any]) -> None:
+    model_state_dict = model.state_dict()
+    _substitute_if_same_shape(model_state_dict, state_dict)
+    model.load_state_dict(model_state_dict)
+
+
 def load_checkpoint(
     checkpoint_path: Path | str,
     model: torch.nn.Module,
@@ -174,37 +200,22 @@ def load_checkpoint(
     checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
     iteration = checkpoint_dict["iteration"]
     learning_rate = checkpoint_dict["learning_rate"]
+
+    # safe load module
+    if hasattr(model, "module"):
+        safe_load(model.module, checkpoint_dict["model"])
+    else:
+        safe_load(model, checkpoint_dict["model"])
+    # safe load optim
     if (
         optimizer is not None
         and not skip_optimizer
         and checkpoint_dict["optimizer"] is not None
     ):
-        try:
-            optimizer.load_state_dict(checkpoint_dict["optimizer"])
-        except Exception as e:
-            LOG.exception(e)
-            LOG.warning("Failed to load optimizer state")
-    saved_state_dict = checkpoint_dict["model"]
-    if hasattr(model, "module"):
-        state_dict = model.module.state_dict()
-    else:
-        state_dict = model.state_dict()
-    new_state_dict = {}
-    for k, v in state_dict.items():
-        try:
-            new_state_dict[k] = saved_state_dict[k]
-            assert saved_state_dict[k].shape == v.shape, (
-                saved_state_dict[k].shape,
-                v.shape,
-            )
-        except Exception as e:
-            LOG.exception(e)
-            LOG.error("%s is not in the checkpoint" % k)
-            new_state_dict[k] = v
-    if hasattr(model, "module"):
-        model.module.load_state_dict(new_state_dict)
-    else:
-        model.load_state_dict(new_state_dict)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            safe_load(optimizer, checkpoint_dict["optimizer"])
+
     LOG.info(f"Loaded checkpoint '{checkpoint_path}' (iteration {iteration})")
     return model, optimizer, learning_rate, iteration
 


Test patch: