"python ./run_csm.py" error: torch._dynamo.exc.RecompileLimitExceeded: cache_size_limit reached

[jeffDebug]

Any ideas and fixes to this error? Thanks!

After cloning the repo: https://github.com/davidbrowne17/csm-streaming, and following setup steps, I then tried to run "python ./run_csm.py" in the repo directory. I got this error. The following is the logging and full stack trace on my terminal.

(.venv) csm-streaming$ python ./run_csm.py
Using device: cuda
Loading CSM-1B model with extreme optimizations for real-time performance...
Model compilation complete. Creating generator...
Starting maximum-intensity warmup sequence...
Optimizing GPU memory allocation...
Creating diverse audio contexts...
Forcing compilation of critical components...
Running final generation with exact same setup as a real request...
Preparing GPU for low-latency generation...
Starting audio generation for: 'This is the final warmup that exactly matches a re...'
First chunk latency: 26059.0ms
Total time: 26.77s
Generated 52 frames (4.16s of audio)
Real-time factor: 6.435x (target: <1.0)

==================================================
AUDIO GENERATION PERFORMANCE METRICS

First chunk latency: 26059.0ms
Total generation time: 27.30s
Audio duration: 4.16s
Real-time factor (RTF): 6.562x (target: <1.0)
Number of chunks: 3
Average chunk size: 1386.7ms
Average inter-chunk latency: 368.7ms
Min/Max inter-chunk latency: 246.2ms / 491.2ms
Chunks per second: 0.11
Output file: warmup_final.wav

Final GPU optimization...
Maximum-intensity warmup complete. First generation should now be MUCH faster.
Generating: Hey how are you doing?
W0521 20:49:28.140000 683560 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:906] [0/8] torch._dynamo hit config.cache_size_limit (8)
W0521 20:49:28.140000 683560 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:906] [0/8] function: 'forward' (/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torchtune/modules/transformer.py:554)
W0521 20:49:28.140000 683560 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:906] [0/8] last reason: 0/0: tensor 'L['mask']' size mismatch at index 1. expected 10, actual 1
W0521 20:49:28.140000 683560 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:906] [0/8] To log all recompilation reasons, use TORCH_LOGS="recompiles".
W0521 20:49:28.140000 683560 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:906] [0/8] To diagnose recompilation issues, see https://pytorch.org/docs/main/torch.compiler_troubleshooting.html.
Traceback (most recent call last):
File "/home/tmp/csm-streaming/./run_csm.py", line 115, in
main()
File "/home/tmp/csm-streaming/./run_csm.py", line 97, in main
audio_tensor = generator.generate(
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/home/tmp/csm-streaming/generator.py", line 441, in generate
sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
File "/home/tmp/csm-streaming/models.py", line 152, in generate_frame
h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask).to(dtype=dtype)
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 574, in _fn
return fn(*args, **kwargs)
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1380, in call
return self._torchdynamo_orig_callable(
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 547, in call
return _compile(
File "/home/tmp/csm-streaming/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 925, in _compile
raise RecompileLimitExceeded(f"{limit_type} reached")
torch._dynamo.exc.RecompileLimitExceeded: cache_size_limit reached

[jeffDebug]

By the way, if I modify the audio generating call to the stream version, in run_csm.py, as the following, it worked fine.

for utterance in conversation:
      print(f"Generating: {utterance['text']}")
      audio_tensor = torch.cat(
          list(generator.generate_stream(
              text=utterance['text'],
              speaker=utterance['speaker_id'],
              context=prompt_segments + generated_segments,
              max_audio_length_ms=10_000,
          )),
          dim=0,
    )
    generated_segments.append(Segment(text=utterance['text'], speaker=utterance['speaker_id'], audio=audio_tensor))

[HarrisonUnifyAI]

I am experiencing the same issue. If I add "stream=True" I can get samples to work. But when I move this into a flask server to serve up an API, it always fails. I'm basically hot-swapping between csm-streaming and the main csm branch

The csm branch works in the Flask API, but it's very slow. The csm-streaming branch works in startup but not in the API call, always failing with:

File "/home/ubuntu/UnifyAI/unifyai/third_party/csm_streaming/generator.py", line 401, in generate
    for chunk in self.generate_stream(text, speaker, context, max_audio_length_ms, temperature, topk):
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 36, in generator_context
    response = gen.send(None)
  File "/home/ubuntu/UnifyAI/unifyai/third_party/csm_streaming/generator.py", line 242, in generate_stream
    sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
  File "/home/ubuntu/UnifyAI/unifyai/third_party/csm_streaming/models.py", line 152, in generate_frame
    h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask).to(dtype=dtype)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 574, in _fn
    return fn(*args, **kwargs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torchtune/modules/transformer.py", line 554, in forward
    def forward(
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 745, in _fn
    return fn(*args, **kwargs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1184, in forward
    return compiled_fn(full_args)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 323, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 126, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 672, in inner_fn
    outs = compiled_fn(args)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 490, in wrapper
    return compiled_fn(runtime_args)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_inductor/output_code.py", line 466, in __call__
    return self.current_callable(inputs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1208, in run
    return compiled_fn(new_inputs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 398, in deferred_cudagraphify
    fn, out = cudagraphify(model, inputs, new_static_input_idxs, *args, **kwargs)
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 420, in cudagraphify
    manager = get_container(device_index).get_tree_manager()
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 341, in get_container
    container_dict = get_obj(local, "tree_manager_containers")
  File "/home/ubuntu/environment/unify_ai/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 336, in get_obj
    assert torch._C._is_key_in_tls(attr_name)

My code basically looks like this:

@app.route('/chat', methods=['POST'])
def chat():
   text = (get text from request)
   speech_gen.generate(speaker_id=0, text=text) # fails here (original csm branch works here)

if __name__ == "__main__":
   (server setup, create speech gen implementation)
   speech_gen.generate(speaker_id=0, text=someText)
   speech_gen.generate(speaker_id=0, text=someText2)
   speech_gen.generate(speaker_id=0, text=someText3)
   speech_gen.generate(speaker_id=0, text=someText4)
   # all of these pass without issue

[HarrisonUnifyAI]

Seems related to this issue (a multithreading issue), not sure if the maintainer of this lib can do anything. The workaround seems to work for now but it's apparently unstable

pytorch/pytorch#123177