llama3.1 405B

Describe the bug
https://wandb.ai/Nemo-automodel/trending-4kseqlen-20251009/runs/t35l7rwa/logs

Steps/Code to reproduce bug
Please see wandb log

Expected behavior

A clear and concise description of what you expected to happen.

Additional context
Stacktrace:

[rank124]:[W1015 18:21:20.649331230 socket.cpp:460] [c10d] waitForInput: poll for socket SocketImpl(fd=93, addr=[nousresearch-hermes-4-405b-sfttp8p-62z4-15-lan.nousresearch-hermes-4-405b-sfttp8p-62z4-job-svc.ws-y5sal9c1.svc.cluster.local]:34008, remote=[nousresearch-hermes-4-405b-sfttp8p-62z4-0-lan.nousresearch-hermes-4-405b-sfttp8p-62z4-job-svc.ws-y5sal9c1.svc.cluster.local]:29400) returned 0, likely a timeout
[rank124]:[W1015 18:21:20.155929192 socket.cpp:485] [c10d] waitForInput: socket SocketImpl(fd=93, addr=[nousresearch-hermes-4-405b-sfttp8p-62z4-15-lan.nousresearch-hermes-4-405b-sfttp8p-62z4-job-svc.ws-y5sal9c1.svc.cluster.local]:34008, remote=[nousresearch-hermes-4-405b-sfttp8p-62z4-0-lan.nousresearch-hermes-4-405b-sfttp8p-62z4-job-svc.ws-y5sal9c1.svc.cluster.local]:29400) timed out after 60000ms
[rank124]: Traceback (most recent call last):
[rank124]:   File "/nemo-workspace/huiyingl/Automodel/examples/llm_finetune/finetune.py", line 33, in <module>
[rank124]:     main()
[rank124]:   File "/nemo-workspace/huiyingl/Automodel/examples/llm_finetune/finetune.py", line 29, in main
[rank124]:     recipe.run_train_validation_loop()
[rank124]:   File "/nemo-workspace/huiyingl/Automodel/nemo_automodel/recipes/llm/train_ft.py", line 890, in run_train_validation_loop
[rank124]:     reporting_loss, grad_norm, tps, num_tokens_in_batch, num_label_tokens = self._run_train_optim_step(
[rank124]:                                                                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank124]:   File "/nemo-workspace/huiyingl/Automodel/nemo_automodel/recipes/llm/train_ft.py", line 995, in _run_train_optim_step
[rank124]:     self._forward_backward_step(
[rank124]:   File "/nemo-workspace/huiyingl/Automodel/nemo_automodel/recipes/llm/train_ft.py", line 938, in _forward_backward_step
[rank124]:     self.pp.info.schedule.step(target=targets, losses=losses, **batch)
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/pipelining/schedules.py", line 1244, in step
[rank124]:     self._step_microbatches(args_split, kwargs_split, targets_split, losses)
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/pipelining/schedules.py", line 1269, in _step_microbatches
[rank124]:     self._initialize_stages(arg_mbs[0], kwarg_mbs[0])
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/pipelining/schedules.py", line 1173, in _initialize_stages
[rank124]:     next_stage_args = stage._prepare_forward_infra(
[rank124]:                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/pipelining/stage.py", line 1448, in _prepare_forward_infra
[rank124]:     outputs = self._shape_inference(args, kwargs)
[rank124]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/pipelining/stage.py", line 1359, in _shape_inference
[rank124]:     dist.recv_object_list(
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/c10d_logger.py", line 81, in wrapper
[rank124]:     return func(*args, **kwargs)
[rank124]:            ^^^^^^^^^^^^^^^^^^^^^
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py", line 3443, in recv_object_list
[rank124]:     rank_sizes = recv(object_sizes_tensor, src=src, group=group, group_src=group_src)
[rank124]:                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/c10d_logger.py", line 81, in wrapper
[rank124]:     return func(*args, **kwargs)
[rank124]:            ^^^^^^^^^^^^^^^^^^^^^
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py", line 2473, in recv
[rank124]:     work = irecv(tensor, src=src, group=group, tag=tag, group_src=group_src)
[rank124]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank124]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py", line 2411, in irecv
[rank124]:     return group.recv([tensor], group_src, tag)
[rank124]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

TORCH_NCCL_TRACE_BUFFER_SIZE

Start with examples/llm_finetune/finetune.py and nemo_automodel/recipes/llm/train_ft.py, then follow the PyTorch distributed pipelining calls into shape inference and recv_object_list. Review the linked W&B log and stack trace to identify the timeout's reproducible conditions; done means the 405B training run completes without this distributed socket timeout.