Full fine-tuning on Alpaca dataset with 4 L40s GPUs fails 8 hours into the training job with index_copy_ about lit-llama HOT 2 CLOSED

[2023-11-03 19:13:17.491399] iter 7998: loss 0.8175, time: 1408.26ms
Validating ...
Traceback (most recent call last):
  File "/root/lit-llama/finetune/full.py", line 225, in <module>
    CLI(main)
  File "/root/.venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 96, in CLI
    return _run_component(components, cfg_init)
  File "/root/.venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 181, in _run_component
    return component(**cfg)
  File "/root/lit-llama/finetune/full.py", line 86, in main
    train(fabric, model, optimizer, train_data, val_data, out_dir)
  File "/root/lit-llama/finetune/full.py", line 131, in train
    val_loss = validate(fabric, model, val_data)
Traceback (most recent call last):
  File "/root/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
  File "/root/lit-llama/finetune/full.py", line 225, in <module>
    return func(*args, **kwargs)
  File "/root/lit-llama/finetune/full.py", line 177, in validate
    output = generate_response(model, instruction)
  File "/root/lit-llama/finetune/full.py", line 152, in generate_response
    output = generate(
  File "/root/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
Traceback (most recent call last):
  File "/root/lit-llama/finetune/full.py", line 225, in <module>
    CLI(main)
  File "/root/.venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 96, in CLI
    return func(*args, **kwargs)
  File "/root/lit-llama/generate.py", line 83, in generate
    return _run_component(components, cfg_init)
    idx = idx.index_copy(0, input_pos, idx_next)
  File "/root/.venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 181, in _run_component
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument source in method wrapper_CUDA_index_copy)
    return component(**cfg)
  File "/root/lit-llama/finetune/full.py", line 86, in main
    train(fabric, model, optimizer, train_data, val_data, out_dir)
  File "/root/lit-llama/finetune/full.py", line 131, in train
    CLI(main)
  File "/root/.venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 96, in CLI
    val_loss = validate(fabric, model, val_data)
  File "/root/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return _run_component(components, cfg_init)
    return func(*args, **kwargs)
  File "/root/.venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 181, in _run_component
  File "/root/lit-llama/finetune/full.py", line 177, in validate
    return component(**cfg)
    output = generate_response(model, instruction)
  File "/root/lit-llama/finetune/full.py", line 86, in main
  File "/root/lit-llama/finetune/full.py", line 152, in generate_response
    train(fabric, model, optimizer, train_data, val_data, out_dir)
  File "/root/lit-llama/finetune/full.py", line 131, in train
    output = generate(
  File "/root/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
    val_loss = validate(fabric, model, val_data)
  File "/root/lit-llama/generate.py", line 83, in generate
  File "/root/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    idx = idx.index_copy(0, input_pos, idx_next)
    return func(*args, **kwargs)
  File "/root/lit-llama/finetune/full.py", line 177, in validate
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:3! (when checking argument for argument source in method wrapper_CUDA_index_copy)
    output = generate_response(model, instruction)
  File "/root/lit-llama/finetune/full.py", line 152, in generate_response
    output = generate(
  File "/root/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/lit-llama/generate.py", line 83, in generate
    idx = idx.index_copy(0, input_pos, idx_next)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:2! (when checking argument for argument source in method wrapper_CUDA_index_copy)
Traceback (most recent call last):
  File "/root/lit-llama/finetune/full.py", line 225, in <module>
    CLI(main)
  File "/root/.venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 96, in CLI
    return _run_component(components, cfg_init)
  File "/root/.venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 181, in _run_component
    return component(**cfg)
  File "/root/lit-llama/finetune/full.py", line 86, in main
    train(fabric, model, optimizer, train_data, val_data, out_dir)
  File "/root/lit-llama/finetune/full.py", line 131, in train
    val_loss = validate(fabric, model, val_data)
  File "/root/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/lit-llama/finetune/full.py", line 177, in validate
    output = generate_response(model, instruction)
  File "/root/lit-llama/finetune/full.py", line 152, in generate_response
    output = generate(
  File "/root/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/lit-llama/generate.py", line 83, in generate
    idx = idx.index_copy(0, input_pos, idx_next)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:1! (when checking argument for argument source in method wrapper_CUDA_index_copy)