ai21labs

/

Jamba-v0.1

like 1.17k

Follow

AI21 185

• transformers Version: I tried 4.40.0 and 4.44.0
• pytorch version: 2.3.1

Code

from transformers import AutoModelForCausalLM, AutoTokenizer, JambaConfig

config = JambaConfig(use_mamba_kernels=False)
model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/Jamba-v0.1",
    torch_dtype=torch.float16,
    config=config,
    attn_implementation="flash_attention_2",
    device_map="auto",
)

tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")

input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]
outputs = model.generate(input_ids, max_new_tokens=216)
print(tokenizer.batch_decode(outputs))

Error Details

AttributeError                            Traceback (most recent call last)
Cell In[13], line 1
----> 1 outputs = model.generate(input_ids, max_new_tokens=216)
      2 print(tokenizer.batch_decode(outputs))
      3 # ["<|startoftext|>In the recent Super Bowl LVIII, the Kansas City Chiefs emerged victorious, defeating the San Francisco 49ers in a thrilling overtime showdown. The game was a nail-biter, with both teams showcasing their skills and determination.\n\nThe Chiefs, led by their star quarterback Patrick Mahomes, displayed their offensive prowess, while the 49ers, led by their strong defense, put up a tough fight. The game went into overtime, with the Chiefs ultimately securing the win with a touchdown.\n\nThe victory marked the Chiefs' second Super Bowl win in four years, solidifying their status as one of the top teams in the NFL. The game was a testament to the skill and talent of both teams, and a thrilling end to the NFL season.\n\nThe Super Bowl is not just about the game itself, but also about the halftime show and the commercials. This year's halftime show featured a star-studded lineup, including Usher, Alicia Keys, and Lil Jon. The show was a spectacle of music and dance, with the performers delivering an energetic and entertaining performance.\n"]

File ~/.local/lib/python3.11/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    112 @functools.wraps(func)
    113 def decorate_context(*args, **kwargs):
    114     with ctx_factory():
--> 115         return func(*args, **kwargs)

File ~/.local/lib/python3.11/site-packages/transformers/generation/utils.py:2024, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   2016     input_ids, model_kwargs = self._expand_inputs_for_generation(
   2017         input_ids=input_ids,
   2018         expand_size=generation_config.num_return_sequences,
   2019         is_encoder_decoder=self.config.is_encoder_decoder,
   2020         **model_kwargs,
   2021     )
   2023     # 13. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
-> 2024     result = self._sample(
   2025         input_ids,
   2026         logits_processor=prepared_logits_processor,
   2027         logits_warper=prepared_logits_warper,
   2028         stopping_criteria=prepared_stopping_criteria,
   2029         generation_config=generation_config,
   2030         synced_gpus=synced_gpus,
   2031         streamer=streamer,
   2032         **model_kwargs,
   2033     )
   2035 elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH):
   2036     # 11. prepare logits warper
   2037     prepared_logits_warper = (
   2038         self._get_logits_warper(generation_config, device=input_ids.device)
   2039         if generation_config.do_sample
   2040         else None
   2041     )

File ~/.local/lib/python3.11/site-packages/transformers/generation/utils.py:2982, in GenerationMixin._sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, logits_warper, **model_kwargs)
   2979 model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
   2981 # forward pass to get next token
-> 2982 outputs = self(**model_inputs, return_dict=True)
   2984 if synced_gpus and this_peer_finished:
   2985     continue  # don't waste resources running the code we don't need

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/.local/lib/python3.11/site-packages/accelerate/hooks.py:166, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    164         output = module._old_forward(*args, **kwargs)
    165 else:
--> 166     output = module._old_forward(*args, **kwargs)
    167 return module._hf_hook.post_forward(module, output)

File ~/.local/lib/python3.11/site-packages/transformers/models/jamba/modeling_jamba.py:1480, in JambaForCausalLM.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, output_router_logits, return_dict, cache_position, num_logits_to_keep)
   1477 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   1479 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> 1480 outputs = self.model(
   1481     input_ids=input_ids,
   1482     attention_mask=attention_mask,
   1483     position_ids=position_ids,
   1484     past_key_values=past_key_values,
   1485     inputs_embeds=inputs_embeds,
   1486     use_cache=use_cache,
   1487     output_attentions=output_attentions,
   1488     output_hidden_states=output_hidden_states,
   1489     output_router_logits=output_router_logits,
   1490     cache_position=cache_position,
   1491     return_dict=return_dict,
   1492 )
   1494 hidden_states = outputs[0]
   1495 if num_logits_to_keep is None:

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/.local/lib/python3.11/site-packages/transformers/models/jamba/modeling_jamba.py:1302, in JambaModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, output_router_logits, return_dict, cache_position)
   1290     layer_outputs = self._gradient_checkpointing_func(
   1291         decoder_layer.__call__,
   1292         hidden_states,
   (...)
   1299         cache_position,
   1300     )
   1301 else:
-> 1302     layer_outputs = decoder_layer(
   1303         hidden_states,
   1304         attention_mask=causal_mask,
   1305         position_ids=position_ids,
   1306         past_key_value=past_key_values,
   1307         output_attentions=output_attentions,
   1308         output_router_logits=output_router_logits,
   1309         use_cache=use_cache,
   1310         cache_position=cache_position,
   1311     )
   1313 hidden_states = layer_outputs[0]
   1315 if output_attentions:

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/.local/lib/python3.11/site-packages/accelerate/hooks.py:161, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    160 def new_forward(module, *args, **kwargs):
--> 161     args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
    162     if module._hf_hook.no_grad:
    163         with torch.no_grad():

File ~/.local/lib/python3.11/site-packages/accelerate/hooks.py:356, in AlignDevicesHook.pre_forward(self, module, *args, **kwargs)
    345             self.tied_pointers_to_remove.add((value.data_ptr(), self.execution_device))
    347         set_module_tensor_to_device(
    348             module,
    349             name,
   (...)
    353             tied_params_map=self.tied_params_map,
    354         )
--> 356 return send_to_device(args, self.execution_device), send_to_device(
    357     kwargs, self.execution_device, skip_keys=self.skip_keys
    358 )

File ~/.local/lib/python3.11/site-packages/accelerate/utils/operations.py:186, in send_to_device(tensor, device, non_blocking, skip_keys)
    183     elif skip_keys is None:
    184         skip_keys = []
    185     return type(tensor)(
--> 186         {
    187             k: t if k in skip_keys else send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys)
    188             for k, t in tensor.items()
    189         }
    190     )
    191 else:
    192     return tensor

File ~/.local/lib/python3.11/site-packages/accelerate/utils/operations.py:187, in <dictcomp>(.0)
    183     elif skip_keys is None:
    184         skip_keys = []
    185     return type(tensor)(
    186         {
--> 187             k: t if k in skip_keys else send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys)
    188             for k, t in tensor.items()
    189         }
    190     )
    191 else:
    192     return tensor

File ~/.local/lib/python3.11/site-packages/accelerate/utils/operations.py:158, in send_to_device(tensor, device, non_blocking, skip_keys)
    156     tensor = tensor.cpu()
    157 try:
--> 158     return tensor.to(device, non_blocking=non_blocking)
    159 except TypeError:  # .to() doesn't accept non_blocking as kwarg
    160     return tensor.to(device)

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1173, in Module.to(self, *args, **kwargs)
   1170         else:
   1171             raise
-> 1173 return self._apply(convert)

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:778, in Module._apply(self, fn, recurse)
    776 def _apply(self, fn, recurse=True):
    777     if recurse:
--> 778         for module in self.children():
    779             module._apply(fn)
    781     def compute_should_use_set_data(tensor, tensor_applied):

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:2319, in Module.children(self)
   2313 def children(self) -> Iterator['Module']:
   2314     r"""Return an iterator over immediate children modules.
   2315 
   2316     Yields:
   2317         Module: a child module
   2318     """
-> 2319     for name, module in self.named_children():
   2320         yield module

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:2337, in Module.named_children(self)
   2323 r"""Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
   2324 
   2325 Yields:
   (...)
   2334 
   2335 """
   2336 memo = set()
-> 2337 for name, module in self._modules.items():
   2338     if module is not None and module not in memo:
   2339         memo.add(module)

File ~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1709, in Module.__getattr__(self, name)
   1707     if name in modules:
   1708         return modules[name]
-> 1709 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

AttributeError: 'HybridMambaAttentionDynamicCache' object has no attribute '_modules'

Any help or pointers would be greatly appreciated!

Thank you!