TP broken

[songhappy]

#2782 broke TP, the previous commit was fine. Error messages look like:

[rank2]:   File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 956, in train
[rank2]:     current_loss.backward()
[rank2]:   File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/_tensor.py", line 648, in backward
[rank2]:     torch.autograd.backward(
[rank2]:   File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/__init__.py", line 353, in backward
[rank2]:     _engine_run_backward(
[rank2]:   File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
[rank2]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank2]: RuntimeError: Function SliceBackward0 returned an invalid gradient at index 0 - got [753, 4096] but expected shape compatible with [763, 4096] 

Steps to reproduce:

 git checkout 9983bbc1dd137f17f951d080f5638050ced432af
 tune run  --nproc_per_node 4 full_finetune_distributed  --config $YOUR_CONFIG\
     seed=123 dataset.packed=True log_level=DEBUG\
    tokenizer.path=$MODEL_PATH/original/tokenizer.model \
    checkpointer.checkpoint_dir=$MODEL_PATH \
    tensor_parallel_dim=2 \
    output_dir=$OUTPUT_PATH/llama3_1/8b_full_tp2 \
    max_steps_per_epoch=$max_steps_per_epoch tokenizer.max_seq_len=512 batch_size=2 

[songhappy]

@nathan-az Would you please help out?

[nathan-az]

Interesting. Two requests:

1. Can you share your full config?
2. Can you try disabling the masking for the loss component? loss.mask_ignored_tokens: False

[songhappy]

with loss.mask_ignored_tokens: False, it runs smoothly on cuda, need further investigation on xpu.
full error log looks like this, you can also reproduce it using previous steps.
`INFO:torchtune.utils._logging: Profiler config after instantiation: {'enabled': False}
0%| | 0/10 [00:00<?, ?it/s]DEBUG:torchtune.utils._logging:Using flex attention for attention computation since a BlockMask was passed in.
1|3|Loss: 1.2586880922317505: 30%|█████████████████████▌ | 3/10 [00:05<00:11, 1.64s/it][rank0]: Traceback (most recent call last):
[rank0]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 1117, in
[rank0]: sys.exit(recipe_main())
[rank0]: File "/home/songhappy/git/torchtune/torchtune/config/_parse.py", line 99, in wrapper
[rank0]: sys.exit(recipe_main(conf))
[rank0]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 1112, in recipe_main
[rank0]: recipe.train()
[rank0]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 956, in train
[rank0]: current_loss.backward()
[rank0]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/_tensor.py", line 648, in backward
[rank0]: torch.autograd.backward(
[rank0]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/init.py", line 353, in backward
[rank0]: _engine_run_backward(
[rank0]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
[rank0]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank0]: RuntimeError: Function SliceBackward0 returned an invalid gradient at index 0 - got [753, 4096] but expected shape compatible with [763, 4096]
[rank3]: Traceback (most recent call last):
[rank3]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 1117, in
[rank3]: sys.exit(recipe_main())
[rank3]: File "/home/songhappy/git/torchtune/torchtune/config/_parse.py", line 99, in wrapper
[rank3]: sys.exit(recipe_main(conf))
[rank3]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 1112, in recipe_main
[rank3]: recipe.train()
[rank3]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 956, in train
[rank3]: current_loss.backward()
[rank3]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/_tensor.py", line 648, in backward
[rank3]: torch.autograd.backward(
[rank3]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/init.py", line 353, in backward
[rank3]: _engine_run_backward(
[rank3]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
[rank3]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank3]: RuntimeError: Function SliceBackward0 returned an invalid gradient at index 0 - got [753, 4096] but expected shape compatible with [763, 4096]
[rank1]: Traceback (most recent call last):
[rank1]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 1117, in
[rank1]: sys.exit(recipe_main())
[rank1]: File "/home/songhappy/git/torchtune/torchtune/config/_parse.py", line 99, in wrapper
[rank1]: sys.exit(recipe_main(conf))
[rank1]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 1112, in recipe_main
[rank1]: recipe.train()
[rank1]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 956, in train
[rank1]: current_loss.backward()
[rank1]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/_tensor.py", line 648, in backward
[rank1]: torch.autograd.backward(
[rank1]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/init.py", line 353, in backward
[rank1]: _engine_run_backward(
[rank1]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
[rank1]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank1]: RuntimeError: Function SliceBackward0 returned an invalid gradient at index 0 - got [753, 4096] but expected shape compatible with [763, 4096]
[rank2]: Traceback (most recent call last):
[rank2]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 1117, in
[rank2]: sys.exit(recipe_main())
[rank2]: File "/home/songhappy/git/torchtune/torchtune/config/_parse.py", line 99, in wrapper
[rank2]: sys.exit(recipe_main(conf))
[rank2]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 1112, in recipe_main
[rank2]: recipe.train()
[rank2]: File "/home/songhappy/git/torchtune/recipes/full_finetune_distributed.py", line 956, in train
[rank2]: current_loss.backward()
[rank2]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/_tensor.py", line 648, in backward
[rank2]: torch.autograd.backward(
[rank2]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/init.py", line 353, in backward
[rank2]: _engine_run_backward(
[rank2]: File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
[rank2]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank2]: RuntimeError: Function SliceBackward0 returned an invalid gradient at index 0 - got [753, 4096] but expected shape compatible with [763, 4096]
1|3|Loss: 1.2586880922317505: 30%|█████████████████████▌ | 3/10 [00:06<00:16, 2.32s/it]
[rank0]:[W714 18:32:10.044397265 ProcessGroupNCCL.cpp:1479] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
[rank1]:[W714 18:32:10.739012361 ProcessGroup.cpp:266] Warning: At the time of process termination, there are still 7 unwaited collective calls. Please review your program to ensure that:

1. c10d_functional.wait_tensor() is invoked on all tensors returned from c10d_functional collective,
2. c10d_functional.wait_tensor() is invoked on all output tensors of async_op=True torch.distributed collective called under with allow_inflight_collective_as_graph_input_ctx():,
   before the output tensors of the collective are used. (function ~WorkRegistry)
   [rank2]:[W714 18:32:10.748043656 ProcessGroup.cpp:266] Warning: At the time of process termination, there are still 7 unwaited collective calls. Please review your program to ensure that:
3. c10d_functional.wait_tensor() is invoked on all tensors returned from c10d_functional collective,
4. c10d_functional.wait_tensor() is invoked on all output tensors of async_op=True torch.distributed collective called under with allow_inflight_collective_as_graph_input_ctx():,
   before the output tensors of the collective are used. (function ~WorkRegistry)
   [rank0]:[W714 18:32:10.749042318 ProcessGroup.cpp:266] Warning: At the time of process termination, there are still 7 unwaited collective calls. Please review your program to ensure that:
5. c10d_functional.wait_tensor() is invoked on all tensors returned from c10d_functional collective,
6. c10d_functional.wait_tensor() is invoked on all output tensors of async_op=True torch.distributed collective called under with allow_inflight_collective_as_graph_input_ctx():,
   before the output tensors of the collective are used. (function ~WorkRegistry)
   [rank3]:[W714 18:32:11.922285798 ProcessGroup.cpp:266] Warning: At the time of process termination, there are still 7 unwaited collective calls. Please review your program to ensure that:
7. c10d_functional.wait_tensor() is invoked on all tensors returned from c10d_functional collective,
8. c10d_functional.wait_tensor() is invoked on all output tensors of async_op=True torch.distributed collective called under with allow_inflight_collective_as_graph_input_ctx():,
   before the output tensors of the collective are used. (function ~WorkRegistry)
   W0714 18:32:11.430000 3456762 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 3456968 closing signal SIGTERM
   W0714 18:32:11.431000 3456762 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 3456970 closing signal SIGTERM
   W0714 18:32:11.431000 3456762 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 3456971 closing signal SIGTERM
   E0714 18:32:11.772000 3456762 site-packages/torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: 1) local_rank: 1 (pid: 3456969) of binary: /home/songhappy/miniforge3/envs/guoqiong-pt-2/bin/python3.10
   Traceback (most recent call last):
   File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/bin/tune", line 8, in
   sys.exit(main())
   File "/home/songhappy/git/torchtune/torchtune/_cli/tune.py", line 52, in main
   parser.run(args)
   File "/home/songhappy/git/torchtune/torchtune/_cli/tune.py", line 46, in run
   args.func(args)
   File "/home/songhappy/git/torchtune/torchtune/_cli/run.py", line 212, in _run_cmd
   self._run_distributed(args, is_builtin=is_builtin)
   File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/init.py", line 355, in wrapper
   return f(*args, **kwargs)
   File "/home/songhappy/git/torchtune/torchtune/_cli/run.py", line 101, in _run_distributed
   run(args)
   File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/distributed/run.py", line 883, in run
   elastic_launch(
   File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 139, in call
   return launch_agent(self._config, self._entrypoint, list(args))
   File "/home/songhappy/miniforge3/envs/guoqiong-pt-2/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 270, in launch_agent
   raise ChildFailedError(
   torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
   ============================================================
   /home/songhappy/git/torchtune/recipes/full_finetune_distributed.py FAILED`

[nathan-az]

I'm glad it's working for you on CUDA. Unfortunately I have no way to test what is going wrong on XPU. I've reviewed the code.

My suspicion was that something odd is happening with the masking, but I just realised your error is during the backward pass, so it's unlikely to be a logical error.

This might be an XPU-specific error. Unfortunately I don't know enough about the inner workings of autograd to speculate or look further into this. Apologies.

[songhappy]

could you please fix the error from "loss.mask_ignored_tokens: True" on cuda which is default? I will investigate on xpu.

[nathan-az]

Sorry I'm confused from your message:

with loss.mask_ignored_tokens: True, it runs smoothly on cuda

Does it run correctly with masking, or does it fail with masking on CUDA?

[nathan-az]

I tested TP + loss masking on CUDA in a similar environment to yours (TP size=2, 4 devices with DP sharding inferred) and had no issues:

output_dir: /tmp/torchtune/llama3_1_8B/full

tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  path: /tmp/Meta-Llama-3.1-8B-Instruct/original/tokenizer.model
  max_seq_len: 1024

# Dataset and Sampler
dataset:
  _component_: torchtune.datasets.alpaca_cleaned_dataset
  packed: True
  split: train[:95%]
seed: null
shuffle: True

tensor_parallel_dim: 2
tensor_parallel_plan:
  _component_: torchtune.models.llama3.base_llama_tp_plan

# Validation
run_val_every_n_steps: null
dataset_val:
  _component_: torchtune.datasets.alpaca_cleaned_dataset
  split: train[95%:]
batch_size_val: ${batch_size}

# Model Arguments
model:
  _component_: torchtune.models.llama3_1.llama3_1_8b

checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_dir: /tmp/Meta-Llama-3.1-8B-Instruct/
  checkpoint_files: [
    model-00001-of-00004.safetensors,
    model-00002-of-00004.safetensors,
    model-00003-of-00004.safetensors,
    model-00004-of-00004.safetensors
  ]
  recipe_checkpoint: null
  output_dir: ${output_dir}
  model_type: LLAMA3
resume_from_checkpoint: False

# Fine-tuning arguments
batch_size: 2
epochs: 1

optimizer:
  _component_: torch.optim.AdamW
  lr: 2e-5
  fused: True
loss:
  _component_: torchtune.modules.loss.LinearCrossEntropyLoss
  mask_ignored_tokens: True
max_steps_per_epoch: 5
clip_grad_norm: null
compile: False  # torch.compile the model + loss, True increases speed + decreases memory
optimizer_in_bwd: False  # True saves memory. Requires gradient_accumulation_steps=1
gradient_accumulation_steps: 1  # Use to increase effective batch size

# Training env
device: cuda

# Memory management
enable_activation_checkpointing: True  # True reduces memory
enable_activation_offloading: False  # True reduces memory
custom_sharded_layers: []

# Reduced precision
dtype: bf16

# Logging
metric_logger:
  _component_: torchtune.training.metric_logging.DiskLogger
  log_dir: ${output_dir}/logs
log_every_n_steps: 1
log_peak_memory_stats: True
log_level: INFO  # DEBUG, WARN, etc.

Not much more I can go on here to "fix" if I can't repro the issue. Let me know if there are any stark differences in my test config.

[songhappy]

my bad, with 'loss.mask_ignored_tokens: False', it runs well, the full error message is from 'loss.mask_ignored_tokens: True', that is the default in loss. I can use 'loss.mask_ignored_tokens: False', as workaround if there is no need to fix error from 'loss.mask_ignored_tokens: True' which is default.

[nathan-az]

No worries. I'm happy to leave this open to see if anybody else encounters this issue, especially on CUDA devices (and hopefully can identify a specific trigger in the config).