Can I train with 3*4090 GPU？

Limited by our resource，I have only 3 4090，How can I use these resources for training?

The log is as follows:

The argument trust_remote_code is to be used with Auto classes. It has no effect here and is ignored.
KBLaMPhi3ForCausalLM has generative capabilities, as prepare_inputs_for_generation is explicitly overwritten. However, it doesn't directly inherit from GenerationMixin. From 👉v4.50👈 onwards, PreTrainedModel will NOT inherit from GenerationMixin, and this model will lose the ability to call generate and other related functions.

• If you're using trust_remote_code=True, you can get rid of this warning by loading the model with an auto class. See https://huggingface.co/docs/transformers/en/model_doc/auto#auto-classes
• If you are the owner of the model architecture code, please modify your model class such that it inherits from GenerationMixin (after PreTrainedModel, otherwise you'll get an exception).
• If you are not the owner of the model architecture class, please contact the model code owner to update it.
  Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:02<00:00, 1.02s/it]
  Some weights of KBLaMPhi3ForCausalLM were not initialized from the model checkpoint at microsoft/Phi-3-mini-4k-instruct and are newly initialized: ['model.layers.0.self_attn.q_proj_new.weight', 'model.layers.1.self_attn.q_proj_new.weight', 'model.layers.10.self_attn.q_proj_new.weight', 'model.layers.11.self_attn.q_proj_new.weight', 'model.layers.12.self_attn.q_proj_new.weight', 'model.layers.13.self_attn.q_proj_new.weight', 'model.layers.14.self_attn.q_proj_new.weight', 'model.layers.15.self_attn.q_proj_new.weight', 'model.layers.16.self_attn.q_proj_new.weight', 'model.layers.17.self_attn.q_proj_new.weight', 'model.layers.18.self_attn.q_proj_new.weight', 'model.layers.19.self_attn.q_proj_new.weight', 'model.layers.2.self_attn.q_proj_new.weight', 'model.layers.20.self_attn.q_proj_new.weight', 'model.layers.21.self_attn.q_proj_new.weight', 'model.layers.22.self_attn.q_proj_new.weight', 'model.layers.23.self_attn.q_proj_new.weight', 'model.layers.24.self_attn.q_proj_new.weight', 'model.layers.25.self_attn.q_proj_new.weight', 'model.layers.26.self_attn.q_proj_new.weight', 'model.layers.27.self_attn.q_proj_new.weight', 'model.layers.28.self_attn.q_proj_new.weight', 'model.layers.29.self_attn.q_proj_new.weight', 'model.layers.3.self_attn.q_proj_new.weight', 'model.layers.30.self_attn.q_proj_new.weight', 'model.layers.31.self_attn.q_proj_new.weight', 'model.layers.4.self_attn.q_proj_new.weight', 'model.layers.5.self_attn.q_proj_new.weight', 'model.layers.6.self_attn.q_proj_new.weight', 'model.layers.7.self_attn.q_proj_new.weight', 'model.layers.8.self_attn.q_proj_new.weight', 'model.layers.9.self_attn.q_proj_new.weight']
  You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
  [02:14:09] INFO Phi3Config { train_local.py:914
  "_attn_implementation_autoset": true,
  "_name_or_path": "microsoft/Phi-3-mini-4k-instruct",
  "architectures": [
  "Phi3ForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "auto_map": {
  "AutoConfig": "microsoft/Phi-3-mini-4k-instruct--configuration_phi3.Phi3Config",
  "AutoModelForCausalLM": "microsoft/Phi-3-mini-4k-instruct--modeling_phi3.Phi3ForCausalLM"
  },
  "bos_token_id": 1,
  "embd_pdrop": 0.0,
  "eos_token_id": 32000,
  "hidden_act": "silu",
  "hidden_size": 3072,
  "initializer_range": 0.02,
  "intermediate_size": 8192,
  "max_position_embeddings": 4096,
  "model_type": "phi3",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 32,
  "original_max_position_embeddings": 4096,
  "pad_token_id": 32000,
  "resid_pdrop": 0.0,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 10000.0,
  "sliding_window": 2047,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.48.0",
  "use_cache": true,
  "vocab_size": 32064
  }

[02:14:13] INFO Model ready 🚀 train_local.py:950
WARNING Detected kernel version 5.4.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or other.py:335
higher.
INFO Optimizer recreated train_local.py:568
INFO Number of trainable parameters: 26,222,592.0 train_local.py:971
INFO Training with 1 GPUs train_local.py:595
INFO Total accumulation steps: 20, Steps per GPU: 20 train_local.py:596
INFO Batch size: 20 train_local.py:597
INFO Effective batch size: 400 train_local.py:598
INFO INPUT IDs SHAPE: torch.Size([20, 69]) train_local.py:631
WARNING You are not running the flash-attention implementation, expect numerical differences. logging.py:328
INFO KB SHAPE: torch.Size([20, 5, 33792]) train_local.py:660
INFO GT: What is the purpose of Foreign Investment Promotion Board?<|end|><|assistant|> The purpose of Foreign Investment Promotion Board is promote foreign investment in train_local.py:661
India..<|end|>
INFO PRED: Q is the value of the Directvestment Lawotion and ( train_local.py:662
<|assistant|> The Foreign of the Investment Promotion Board ( to and investment in a. The
[02:14:18] INFO step: 0, loss: 3.564134585857391 train_local.py:704
INFO start_step: 0 train_local.py:709
INFO save_period: 300 train_local.py:710
[02:14:22] INFO step: 1, loss: 3.410089921951294 train_local.py:704
INFO start_step: 0 train_local.py:709
INFO save_period: 300 train_local.py:710
[02:14:26] INFO step: 2, loss: 3.43630336523056 train_local.py:704
INFO start_step: 0 train_local.py:709
INFO save_period: 300 train_local.py:710
[02:14:31] INFO step: 3, loss: 3.3972960829734804 train_local.py:704
INFO start_step: 0 train_local.py:709
INFO save_period: 300 train_local.py:710
[02:14:35] INFO step: 4, loss: 3.275339734554291 train_local.py:704
INFO start_step: 0 train_local.py:709
INFO save_period: 300 train_local.py:710
[02:14:39] INFO step: 5, loss: 3.3063130140304566 train_local.py:704
INFO start_step: 0 train_local.py:709
INFO save_period: 300 train_local.py:710
⠼ Training ╸━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0% Loss: 3.3063 2:23:19
Traceback (most recent call last):
File "/home/vip/gpt/rag/KBLaM/experiments/train_local.py", line 988, in
main()
File "/home/vip/gpt/rag/KBLaM/experiments/train_local.py", line 973, in main
trainer.train(
File "/home/vip/gpt/rag/KBLaM/experiments/train_local.py", line 643, in train
out = self.model(
^^^^^^^^^^^
File "/home/vip/anaconda3/envs/kblam/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/anaconda3/envs/kblam/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/gpt/rag/KBLaM/src/kblam/models/phi3_model.py", line 1090, in forward
outputs = self.model(
^^^^^^^^^^^
File "/home/vip/anaconda3/envs/kblam/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/anaconda3/envs/kblam/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/gpt/rag/KBLaM/src/kblam/models/phi3_model.py", line 932, in forward
layer_outputs = decoder_layer(
^^^^^^^^^^^^^^
File "/home/vip/anaconda3/envs/kblam/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/anaconda3/envs/kblam/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/gpt/rag/KBLaM/src/kblam/models/phi3_model.py", line 630, in forward
hidden_states = self.mlp(hidden_states)
^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/anaconda3/envs/kblam/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/anaconda3/envs/kblam/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vip/gpt/rag/KBLaM/src/kblam/models/phi3_model.py", line 256, in forward
up_states = up_states * self.activation_fn(gate)
~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 46.00 MiB. GPU 0 has a total capacity of 23.53 GiB of which 28.50 MiB is free. Including non-PyTorch memory, this process has 23.46 GiB memory in use. Of the allocated memory 21.56 GiB is allocated by PyTorch, and 1.47 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Review the training entry point referenced in train_local.py around lines 914 and 950, then inspect how the current run initializes the model and optimizer. Reproduce the reported setup with three GPUs and determine whether multi-GPU training is supported. Done should be a documented, reproducible training procedure or a clearly scoped limitation.