FROM nvidia/cuda:11.7.0-cudnn8-devel-ubuntu22.04

RUN apt update
RUN apt install git -y 
RUN apt install wget -y 
RUN apt install python3 python3-pip -y



# Install dependencies (one-by-one for better caching)
#RUN pip install --upgrade pip
RUN pip install torch
RUN pip install transformers
RUN pip install datasets
RUN pip install evaluate
RUN pip install xformers
RUN pip install wandb
RUN pip install peft 
RUN pip install trl 
RUN pip install scipy 
RUN pip install accelerate 
RUN pip install scikit-learn
RUN pip install pandas 
RUN pip install bleurt@https://github.com/google-research/bleurt/archive/b610120347ef22b494b6d69b4316e303f5932516.zip#egg=bleurt

RUN git clone https://github.com/EleutherAI/lm-evaluation-harness
RUN pip install -e lm-evaluation-harness

RUN git clone https://github.com/timdettmers/bitsandbytes.git
# CUDA_VERSIONS in {110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 120}
# make argument in {cuda110, cuda11x, cuda12x}
# if you do not know what CUDA you have, try looking at the output of: python -m bitsandbytes
ENV CUDA_VERSION=117
RUN cd bitsandbytes && git checkout ac5550a0238286377ee3f58a85aeba1c40493e17
RUN cd bitsandbytes && make cuda11x
RUN cd bitsandbytes && python3 setup.py install
#RUN pip install bitsandbytes
#RUN python3 check_bnb_install.py

# Init wandb
#COPY ./wandb /wandb
ENV WANDB_CONFIG_DIR=/wandb

ENV HF_DATASETS_CACHE="/hf_cache/datasets"
ENV HUGGINGFACE_HUB_CACHE="/hf_cache/hub"

# Copy the code
COPY . /code

# Set the working directory
WORKDIR /code

# Install a useful helper to check bitsandbytes installation. Only works at runtime.
RUN wget https://gist.githubusercontent.com/TimDettmers/1f5188c6ee6ed69d211b7fe4e381e713/raw/4d17c3d09ccdb57e9ab7eca0171f2ace6e4d2858/check_bnb_install.py 

---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-23-3435b262f1ae> in <module>
----> 1 trainer.train()

~/.local/lib/python3.8/site-packages/transformers/trainer.py in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1643             self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size
   1644         )
-> 1645         return inner_training_loop(
   1646             args=args,
   1647             resume_from_checkpoint=resume_from_checkpoint,

~/.local/lib/python3.8/site-packages/transformers/trainer.py in _inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1936 
   1937                 with self.accelerator.accumulate(model):
-> 1938                     tr_loss_step = self.training_step(model, inputs)
   1939 
   1940                 if (

~/.local/lib/python3.8/site-packages/transformers/trainer.py in training_step(self, model, inputs)
   2757 
   2758         with self.compute_loss_context_manager():
-> 2759             loss = self.compute_loss(model, inputs)
   2760 
   2761         if self.args.n_gpu > 1:

~/.local/lib/python3.8/site-packages/transformers/trainer.py in compute_loss(self, model, inputs, return_outputs)
   2782         else:
   2783             labels = None
-> 2784         outputs = model(**inputs)
   2785         # Save past state if it exists
   2786         # TODO: this needs to be fixed and made cleaner later.

/usr/lib/python3/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

~/.local/lib/python3.8/site-packages/accelerate/utils/operations.py in forward(*args, **kwargs)
    579 
    580     def forward(*args, **kwargs):
--> 581         return model_forward(*args, **kwargs)
    582 
    583     # To act like a decorator so that it can be popped when doing `extract_model_from_parallel`

~/.local/lib/python3.8/site-packages/accelerate/utils/operations.py in __call__(self, *args, **kwargs)
    567 
    568     def __call__(self, *args, **kwargs):
--> 569         return convert_to_fp32(self.model_forward(*args, **kwargs))
    570 
    571     def __getstate__(self):

/usr/lib/python3/dist-packages/torch/amp/autocast_mode.py in decorate_autocast(*args, **kwargs)
     12     def decorate_autocast(*args, **kwargs):
     13         with autocast_instance:
---> 14             return func(*args, **kwargs)
     15     decorate_autocast.__script_unsupported = '@autocast() decorator is not supported in script mode'  # type: ignore[attr-defined]
     16     return decorate_autocast

~/.local/lib/python3.8/site-packages/peft/peft_model.py in forward(self, *args, **kwargs)
    413         Forward pass of the model.
    414         """
--> 415         return self.get_base_model()(*args, **kwargs)
    416 
    417     def _get_base_model_class(self, is_prompt_tuning=False):

/usr/lib/python3/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

~/.local/lib/python3.8/site-packages/accelerate/hooks.py in new_forward(*args, **kwargs)
    163                 output = old_forward(*args, **kwargs)
    164         else:
--> 165             output = old_forward(*args, **kwargs)
    166         return module._hf_hook.post_forward(module, output)
    167 

~/.local/lib/python3.8/site-packages/transformers/models/whisper/modeling_whisper.py in forward(self, input_features, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, decoder_inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1417                 )
   1418 
-> 1419         outputs = self.model(
   1420             input_features,
   1421             attention_mask=attention_mask,

/usr/lib/python3/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

~/.local/lib/python3.8/site-packages/accelerate/hooks.py in new_forward(*args, **kwargs)
    163                 output = old_forward(*args, **kwargs)
    164         else:
--> 165             output = old_forward(*args, **kwargs)
    166         return module._hf_hook.post_forward(module, output)
    167 

~/.local/lib/python3.8/site-packages/transformers/models/whisper/modeling_whisper.py in forward(self, input_features, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, decoder_inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
   1266             input_features = self._mask_input_features(input_features, attention_mask=attention_mask)
   1267 
-> 1268             encoder_outputs = self.encoder(
   1269                 input_features,
   1270                 head_mask=head_mask,

/usr/lib/python3/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

~/.local/lib/python3.8/site-packages/accelerate/hooks.py in new_forward(*args, **kwargs)
    163                 output = old_forward(*args, **kwargs)
    164         else:
--> 165             output = old_forward(*args, **kwargs)
    166         return module._hf_hook.post_forward(module, output)
    167 

~/.local/lib/python3.8/site-packages/transformers/models/whisper/modeling_whisper.py in forward(self, input_features, attention_mask, head_mask, output_attentions, output_hidden_states, return_dict)
    854                         return custom_forward
    855 
--> 856                     layer_outputs = torch.utils.checkpoint.checkpoint(
    857                         create_custom_forward(encoder_layer),
    858                         hidden_states,

/usr/lib/python3/dist-packages/torch/utils/checkpoint.py in checkpoint(function, use_reentrant, *args, **kwargs)
    247 
    248     if use_reentrant:
--> 249         return CheckpointFunction.apply(function, preserve, *args)
    250     else:
    251         return _checkpoint_without_reentrant(

/usr/lib/python3/dist-packages/torch/autograd/function.py in apply(cls, *args, **kwargs)
    504             # See NOTE: [functorch vjp and autograd interaction]
    505             args = _functorch.utils.unwrap_dead_wrappers(args)
--> 506             return super().apply(*args, **kwargs)  # type: ignore[misc]
    507 
    508         if cls.setup_context == _SingleLevelFunction.setup_context:

/usr/lib/python3/dist-packages/torch/utils/checkpoint.py in forward(ctx, run_function, preserve_rng_state, *args)
    105 
    106         with torch.no_grad():
--> 107             outputs = run_function(*args)
    108         return outputs
    109 

~/.local/lib/python3.8/site-packages/transformers/models/whisper/modeling_whisper.py in custom_forward(*inputs)
    850                     def create_custom_forward(module):
    851                         def custom_forward(*inputs):
--> 852                             return module(*inputs, output_attentions)
    853 
    854                         return custom_forward

/usr/lib/python3/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

~/.local/lib/python3.8/site-packages/accelerate/hooks.py in new_forward(*args, **kwargs)
    163                 output = old_forward(*args, **kwargs)
    164         else:
--> 165             output = old_forward(*args, **kwargs)
    166         return module._hf_hook.post_forward(module, output)
    167 

~/.local/lib/python3.8/site-packages/transformers/models/whisper/modeling_whisper.py in forward(self, hidden_states, attention_mask, layer_head_mask, output_attentions)
    429         residual = hidden_states
    430         hidden_states = self.self_attn_layer_norm(hidden_states)
--> 431         hidden_states, attn_weights, _ = self.self_attn(
    432             hidden_states=hidden_states,
    433             attention_mask=attention_mask,

/usr/lib/python3/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

~/.local/lib/python3.8/site-packages/accelerate/hooks.py in new_forward(*args, **kwargs)
    163                 output = old_forward(*args, **kwargs)
    164         else:
--> 165             output = old_forward(*args, **kwargs)
    166         return module._hf_hook.post_forward(module, output)
    167 

~/.local/lib/python3.8/site-packages/transformers/models/whisper/modeling_whisper.py in forward(self, hidden_states, key_value_states, past_key_value, attention_mask, layer_head_mask, output_attentions)
    288 
    289         # get query proj
--> 290         query_states = self.q_proj(hidden_states) * self.scaling
    291         # get key, value proj
    292         # `past_key_value[0].shape[2] == key_value_states.shape[1]`

/usr/lib/python3/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

~/.local/lib/python3.8/site-packages/peft/tuners/lora.py in forward(self, x)
   1052 
   1053         def forward(self, x: torch.Tensor):
-> 1054             result = super().forward(x)
   1055 
   1056             if self.disable_adapters or self.active_adapter not in self.lora_A.keys():

~/.local/lib/python3.8/site-packages/bitsandbytes/nn/modules.py in forward(self, x)
    412             self.bias.data = self.bias.data.to(x.dtype)
    413 
--> 414         out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state)
    415 
    416         if not self.state.has_fp16_weights:

~/.local/lib/python3.8/site-packages/bitsandbytes/autograd/_functions.py in matmul(A, B, out, state, threshold, bias)
    561     if threshold > 0.0:
    562         state.threshold = threshold
--> 563     return MatMul8bitLt.apply(A, B, out, bias, state)
    564 
    565 

/usr/lib/python3/dist-packages/torch/autograd/function.py in apply(cls, *args, **kwargs)
    504             # See NOTE: [functorch vjp and autograd interaction]
    505             args = _functorch.utils.unwrap_dead_wrappers(args)
--> 506             return super().apply(*args, **kwargs)  # type: ignore[misc]
    507 
    508         if cls.setup_context == _SingleLevelFunction.setup_context:

~/.local/lib/python3.8/site-packages/bitsandbytes/autograd/_functions.py in forward(ctx, A, B, out, bias, state)
    399         if using_igemmlt:
    400             C32A, SA = F.transform(CA, "col32")
--> 401             out32, Sout32 = F.igemmlt(C32A, state.CxB, SA, state.SB)
    402             if bias is None or bias.dtype == torch.float16:
    403                 # we apply the fused bias here

~/.local/lib/python3.8/site-packages/bitsandbytes/functional.py in igemmlt(A, B, SA, SB, out, Sout, dtype)
   1790     if has_error == 1:
   1791         print(f'A: {shapeA}, B: {shapeB}, C: {Sout[0]}; (lda, ldb, ldc): {(lda, ldb, ldc)}; (m, n, k): {(m, n, k)}')
-> 1792         raise Exception('cublasLt ran into an error!')
   1793 
   1794     torch.cuda.set_device(prev_device)

Exception: cublasLt ran into an error!