Vansh Chugh commited on
Commit
df3d717
·
1 Parent(s): 77b301e

disable memory-efficient attention, xformers has no blackwell kernel

Browse files
Files changed (1) hide show
  1. src/audiocraft/models/loaders.py +4 -4
src/audiocraft/models/loaders.py CHANGED
@@ -135,10 +135,10 @@ def load_lm_model_magnet(file_or_url_or_id: tp.Union[Path, str], compression_mod
135
  cfg.transformer_lm.segment_duration = cfg.dataset.segment_duration
136
  cfg.transformer_lm.span_len = cfg.masking.span_len
137
 
138
- # MAGNeT models v1 support only xformers backend.
139
- from src.audiocraft.modules.transformer import set_efficient_attention_backend
140
- if cfg.transformer_lm.memory_efficient:
141
- set_efficient_attention_backend("xformers")
142
 
143
  model = builders.get_lm_model(cfg)
144
  model.load_state_dict(pkg['best_state'])
 
135
  cfg.transformer_lm.segment_duration = cfg.dataset.segment_duration
136
  cfg.transformer_lm.span_len = cfg.masking.span_len
137
 
138
+ # Deviation: xformers has no kernel for ZeroGPU's current hardware
139
+ # (Blackwell, capability 12.0 > xformers' 9.0 cap). Disable memory_efficient
140
+ # rather than switch to 'torch' backend, which ignores custom attn_mask.
141
+ cfg.transformer_lm.memory_efficient = False
142
 
143
  model = builders.get_lm_model(cfg)
144
  model.load_state_dict(pkg['best_state'])