Repository: pytorch
Issue: FSDP1 SHARD_GRAD_OP broken after torch upgrade to 2.4.1 and flash_attn upgrade
Labels: oncall: distributed,triaged,module: fsdp,oncall: pt2,pt2d-triage-nov2024

### 🐛 Describe the bug

Problem seems to be as torch setStorage run. If I recall, the code used to work on torch 2.2 . If I disable FSDP, and the DDP Optimizer, everything works fine. Version fails on 2.4.. I tried setting dynamic to False to fix some issues with Flash Attention complaining a lot and got the same issue.

### Error logs
```
[rank7]:W1029 15:41:39.251000 140228203374400 torch/_dynamo/variables/tensor.py:715] [0/0] Graph break from `Tensor.item()`, consider setting:
[rank7]:W1029 15:41:39.251000 140228203374400 torch/_dynamo/variables/tensor.py:715] [0/0]     torch._dynamo.config.capture_scalar_outputs = True
[rank7]:W1029 15:41:39.251000 140228203374400 torch/_dynamo/variables/tensor.py:715] [0/0] or:
[rank7]:W1029 15:41:39.251000 140228203374400 torch/_dynamo/variables/tensor.py:715] [0/0]     env TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1
[rank7]:W1029 15:41:39.251000 140228203374400 torch/_dynamo/variables/tensor.py:715] [0/0] to include these operations in the captured graph.
[rank7]:W1029 15:41:39.251000 140228203374400 torch/_dynamo/variables/tensor.py:715] [0/0]
[rank7]: [rank7]: [rank7]: ╭───────────────────── Traceback (most recent call last) ──────────────────────╮
[rank7]: [rank7]: [rank7]: │ /llmlib/scripts/train_mosaic_bert.py:287 in <module>                         │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   284 │   cli_cfg = om.from_cli(args_list)                                   │
[rank7]: [rank7]: [rank7]: │   285 │   cfg = om.merge(yaml_cfg, cli_cfg)                                  │
[rank7]: [rank7]: [rank7]: │   286 │   cfg = cast(DictConfig, cfg)  # for type checking                   │
[rank7]: [rank7]: [rank7]: │ ❱ 287 │   main(cfg)                                                          │
[rank7]: [rank7]: [rank7]: │   288                                                                        │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /llmlib/scripts/train_mosaic_bert.py:274 in main                             │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   271 │                                                                      │
[rank7]: [rank7]: [rank7]: │   272 │   if do_train:                                                       │
[rank7]: [rank7]: [rank7]: │   273 │   │   print('Starting training...')                                  │
[rank7]: [rank7]: [rank7]: │ ❱ 274 │   │   trainer.fit()                                                  │
[rank7]: [rank7]: [rank7]: │   275 │                                                                      │
[rank7]: [rank7]: [rank7]: │   276 │   if return_trainer:                                                 │
[rank7]: [rank7]: [rank7]: │   277 │   │   return trainer                                                 │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/trainer/trainer.py:2467 in fit       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   2464 │   │   │   self.state.scaler = ClosureGradScaler() if self._use_clos │
[rank7]: [rank7]: [rank7]: │   2465 │   │                                                                 │
[rank7]: [rank7]: [rank7]: │   2466 │   │   self.first_batch_complete = False                             │
[rank7]: [rank7]: [rank7]: │ ❱ 2467 │   │   self._train_loop()                                            │
[rank7]: [rank7]: [rank7]: │   2468 │   │                                                                 │
[rank7]: [rank7]: [rank7]: │   2469 │   │   # Zero gradients at the end of fit so same model/optimizer ca │
[rank7]: [rank7]: [rank7]: │   2470 │   │   # with checkpoint loading. See https://github.com/pytorch/pyt │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/trainer/trainer.py:2687 in           │
[rank7]: [rank7]: [rank7]: │ _train_loop                                                                  │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   2684 │   │   │   │   │   self.logger.log_metrics({'time/token': self.state │
[rank7]: [rank7]: [rank7]: │   2685 │   │   │   │   │   self.logger.log_metrics({'time/token_in_epoch': s │
[rank7]: [rank7]: [rank7]: │   2686 │   │   │   │                                                         │
[rank7]: [rank7]: [rank7]: │ ❱ 2687 │   │   │   │   total_loss_dict = self._train_batch(use_grad_scaling) │
[rank7]: [rank7]: [rank7]: │   2688 │   │   │   │                                                         │
[rank7]: [rank7]: [rank7]: │   2689 │   │   │   │   if use_grad_scaling:                                  │
[rank7]: [rank7]: [rank7]: │   2690 │   │   │   │   │   self.state.scaler.update()                        │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/trainer/trainer.py:2907 in           │
[rank7]: [rank7]: [rank7]: │ _train_batch                                                                 │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   2904 │   │   │   │   │   │   │   │   **kwargs: self._train_microbatches(mi │
[rank7]: [rank7]: [rank7]: │   2905 │   │   │   │   │   │   │   )                                         │
[rank7]: [rank7]: [rank7]: │   2906 │   │   │   │   │   │   else:                                         │
[rank7]: [rank7]: [rank7]: │ ❱ 2907 │   │   │   │   │   │   │   optimizer.step(                           │
[rank7]: [rank7]: [rank7]: │   2908 │   │   │   │   │   │   │   │   closure=lambda loss_dict=total_loss_d │
[rank7]: [rank7]: [rank7]: │   2909 │   │   │   │   │   │   │   │   **kwargs: self._train_microbatches(mi │
[rank7]: [rank7]: [rank7]: │   2910 │   │   │   │   │   │   │   )                                         │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/optim/lr_scheduler.py:130 in wrapper    │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    127 │   │   │   │   def wrapper(*args, **kwargs):                         │
[rank7]: [rank7]: [rank7]: │    128 │   │   │   │   │   opt = opt_ref()                                   │
[rank7]: [rank7]: [rank7]: │    129 │   │   │   │   │   opt._opt_called = True  # type: ignore[union-attr │
[rank7]: [rank7]: [rank7]: │ ❱  130 │   │   │   │   │   return func.__get__(opt, opt.__class__)(*args, ** │
[rank7]: [rank7]: [rank7]: │    131 │   │   │   │                                                         │
[rank7]: [rank7]: [rank7]: │    132 │   │   │   │   wrapper._wrapped_by_lr_sched = True  # type: ignore[a │
[rank7]: [rank7]: [rank7]: │    133 │   │   │   │   return wrapper                                        │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/optim/optimizer.py:484 in wrapper       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    481 │   │   │   │   │   │   │   │   f"{func} must return None or a tuple  │
[rank7]: [rank7]: [rank7]: │    482 │   │   │   │   │   │   │   )                                         │
[rank7]: [rank7]: [rank7]: │    483 │   │   │   │                                                         │
[rank7]: [rank7]: [rank7]: │ ❱  484 │   │   │   │   out = func(*args, **kwargs)                           │
[rank7]: [rank7]: [rank7]: │    485 │   │   │   │   self._optimizer_step_code()                           │
[rank7]: [rank7]: [rank7]: │    486 │   │   │   │                                                         │
[rank7]: [rank7]: [rank7]: │    487 │   │   │   │   # call optimizer step post hooks                      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/utils/_contextlib.py:116 in             │
[rank7]: [rank7]: [rank7]: │ decorate_context                                                             │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   113 │   @functools.wraps(func)                                             │
[rank7]: [rank7]: [rank7]: │   114 │   def decorate_context(*args, **kwargs):                             │
[rank7]: [rank7]: [rank7]: │   115 │   │   with ctx_factory():                                            │
[rank7]: [rank7]: [rank7]: │ ❱ 116 │   │   │   return func(*args, **kwargs)                               │
[rank7]: [rank7]: [rank7]: │   117 │                                                                      │
[rank7]: [rank7]: [rank7]: │   118 │   return decorate_context                                            │
[rank7]: [rank7]: [rank7]: │   119                                                                        │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/optim/decoupled_weight_decay.py:308  │
[rank7]: [rank7]: [rank7]: │ in step                                                                      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   305 │   │   loss = None                                                    │
[rank7]: [rank7]: [rank7]: │   306 │   │   if closure is not None:                                        │
[rank7]: [rank7]: [rank7]: │   307 │   │   │   with torch.enable_grad():                                  │
[rank7]: [rank7]: [rank7]: │ ❱ 308 │   │   │   │   loss = closure()                                       │
[rank7]: [rank7]: [rank7]: │   309 │   │                                                                  │
[rank7]: [rank7]: [rank7]: │   310 │   │   for group in self.param_groups:                                │
[rank7]: [rank7]: [rank7]: │   311 │   │   │   params_with_grad = []                                      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/trainer/trainer.py:2909 in <lambda>  │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   2906 │   │   │   │   │   │   else:                                         │
[rank7]: [rank7]: [rank7]: │   2907 │   │   │   │   │   │   │   optimizer.step(                           │
[rank7]: [rank7]: [rank7]: │   2908 │   │   │   │   │   │   │   │   closure=lambda loss_dict=total_loss_d │
[rank7]: [rank7]: [rank7]: │ ❱ 2909 │   │   │   │   │   │   │   │   **kwargs: self._train_microbatches(mi │
[rank7]: [rank7]: [rank7]: │   2910 │   │   │   │   │   │   │   )                                         │
[rank7]: [rank7]: [rank7]: │   2911 │   │   │   │   else:                                                 │
[rank7]: [rank7]: [rank7]: │   2912 │   │   │   │   │   self._train_microbatches(microbatches, total_loss │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/trainer/trainer.py:3078 in           │
[rank7]: [rank7]: [rank7]: │ _train_microbatches                                                          │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   3075 │   │   │   for microbatch_idx, self.state.batch in enumerate(microba │
[rank7]: [rank7]: [rank7]: │   3076 │   │   │   │   self.state.batch = self.state.device.batch_to_device( │
[rank7]: [rank7]: [rank7]: │   3077 │   │   │   │   is_final_microbatch = microbatch_idx + 1 == len(micro │
[rank7]: [rank7]: [rank7]: │ ❱ 3078 │   │   │   │   microbatch_loss_dict = self._train_microbatch(use_gra │
[rank7]: [rank7]: [rank7]: │   3079 │   │   │   │                                                         │
[rank7]: [rank7]: [rank7]: │   3080 │   │   │   │   # Aggregate each loss in microbatch_loss_dict into to │
[rank7]: [rank7]: [rank7]: │   3081 │   │   │   │   for k, microbatch_loss in microbatch_loss_dict.items( │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/trainer/trainer.py:3154 in           │
[rank7]: [rank7]: [rank7]: │ _train_microbatch                                                            │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   3151 │   │   │   │   self.state.precision_config,                          │
[rank7]: [rank7]: [rank7]: │   3152 │   │   │   │   self.state.deepspeed_enabled,                         │
[rank7]: [rank7]: [rank7]: │   3153 │   │   │   ):                                                        │
[rank7]: [rank7]: [rank7]: │ ❱ 3154 │   │   │   │   self.state.outputs = self.state.model(self.state.batc │
[rank7]: [rank7]: [rank7]: │   3155 │   │   │                                                             │
[rank7]: [rank7]: [rank7]: │   3156 │   │   │   self.engine.run_event(Event.AFTER_FORWARD)                │
[rank7]: [rank7]: [rank7]: │   3157                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/nn/modules/module.py:1553 in            │
[rank7]: [rank7]: [rank7]: │ _wrapped_call_impl                                                           │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   1550 │   │   if self._compiled_call_impl is not None:                      │
[rank7]: [rank7]: [rank7]: │   1551 │   │   │   return self._compiled_call_impl(*args, **kwargs)  # type: │
[rank7]: [rank7]: [rank7]: │   1552 │   │   else:                                                         │
[rank7]: [rank7]: [rank7]: │ ❱ 1553 │   │   │   return self._call_impl(*args, **kwargs)                   │
[rank7]: [rank7]: [rank7]: │   1554 │                                                                     │
[rank7]: [rank7]: [rank7]: │   1555 │   def _call_impl(self, *args, **kwargs):                            │
[rank7]: [rank7]: [rank7]: │   1556 │   │   forward_call = (self._slow_forward if torch._C._get_tracing_s │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/nn/modules/module.py:1562 in _call_impl │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   1559 │   │   if not (self._backward_hooks or self._backward_pre_hooks or s │
[rank7]: [rank7]: [rank7]: │   1560 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hoo │
[rank7]: [rank7]: [rank7]: │   1561 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks │
[rank7]: [rank7]: [rank7]: │ ❱ 1562 │   │   │   return forward_call(*args, **kwargs)                      │
[rank7]: [rank7]: [rank7]: │   1563 │   │                                                                 │
[rank7]: [rank7]: [rank7]: │   1564 │   │   try:                                                          │
[rank7]: [rank7]: [rank7]: │   1565 │   │   │   result = None                                             │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_dynamo/eval_frame.py:433 in _fn        │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    430 │   │   │   )                                                         │
[rank7]: [rank7]: [rank7]: │    431 │   │   │                                                             │
[rank7]: [rank7]: [rank7]: │    432 │   │   │   try:                                                      │
[rank7]: [rank7]: [rank7]: │ ❱  433 │   │   │   │   return fn(*args, **kwargs)                            │
[rank7]: [rank7]: [rank7]: │    434 │   │   │   finally:                                                  │
[rank7]: [rank7]: [rank7]: │    435 │   │   │   │   # Restore the dynamic layer stack depth if necessary. │
[rank7]: [rank7]: [rank7]: │    436 │   │   │   │   torch._C._functorch.pop_dynamic_layer_stack_and_undo_ │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/models/huggingface.py:488 in forward │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   485 │   │   if isinstance(batch, Mapping):                                 │
[rank7]: [rank7]: [rank7]: │   486 │   │   │   # Further input validation is left to the huggingface forw │
[rank7]: [rank7]: [rank7]: │   487 │   │   │   batch = {k: v for k, v in batch.items() if k in self.model │
[rank7]: [rank7]: [rank7]: │ ❱ 488 │   │   │   output = self.model(**batch)  # type: ignore (thirdparty)  │
[rank7]: [rank7]: [rank7]: │   489 │   │   else:                                                          │
[rank7]: [rank7]: [rank7]: │   490 │   │   │   raise ValueError(                                          │
[rank7]: [rank7]: [rank7]: │   491 │   │   │   │   'Unexpected batch type. Expected a dictionary with key │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/nn/modules/module.py:1553 in            │
[rank7]: [rank7]: [rank7]: │ _wrapped_call_impl                                                           │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   1550 │   │   if self._compiled_call_impl is not None:                      │
[rank7]: [rank7]: [rank7]: │   1551 │   │   │   return self._compiled_call_impl(*args, **kwargs)  # type: │
[rank7]: [rank7]: [rank7]: │   1552 │   │   else:                                                         │
[rank7]: [rank7]: [rank7]: │ ❱ 1553 │   │   │   return self._call_impl(*args, **kwargs)                   │
[rank7]: [rank7]: [rank7]: │   1554 │                                                                     │
[rank7]: [rank7]: [rank7]: │   1555 │   def _call_impl(self, *args, **kwargs):                            │
[rank7]: [rank7]: [rank7]: │   1556 │   │   forward_call = (self._slow_forward if torch._C._get_tracing_s │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/nn/modules/module.py:1603 in _call_impl │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   1600 │   │   │   │   bw_hook = hooks.BackwardHook(self, full_backward_hook │
[rank7]: [rank7]: [rank7]: │   1601 │   │   │   │   args = bw_hook.setup_input_hook(args)                 │
[rank7]: [rank7]: [rank7]: │   1602 │   │   │                                                             │
[rank7]: [rank7]: [rank7]: │ ❱ 1603 │   │   │   result = forward_call(*args, **kwargs)                    │
[rank7]: [rank7]: [rank7]: │   1604 │   │   │   if _global_forward_hooks or self._forward_hooks:          │
[rank7]: [rank7]: [rank7]: │   1605 │   │   │   │   for hook_id, hook in (                                │
[rank7]: [rank7]: [rank7]: │   1606 │   │   │   │   │   *_global_forward_hooks.items(),                   │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/distributed/fsdp/fully_sharded_data_par │
[rank7]: [rank7]: [rank7]: │ allel.py:849 in forward                                                      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    846 │   │   ):                                                            │
[rank7]: [rank7]: [rank7]: │    847 │   │   │   args, kwargs = _root_pre_forward(self, self, args, kwargs │
[rank7]: [rank7]: [rank7]: │    848 │   │   │   unused = None                                             │
[rank7]: [rank7]: [rank7]: │ ❱  849 │   │   │   args, kwargs = _pre_forward(                              │
[rank7]: [rank7]: [rank7]: │    850 │   │   │   │   self,                                                 │
[rank7]: [rank7]: [rank7]: │    851 │   │   │   │   handle,                                               │
[rank7]: [rank7]: [rank7]: │    852 │   │   │   │   _pre_forward_unshard,                                 │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/distributed/fsdp/_runtime_utils.py:381  │
[rank7]: [rank7]: [rank7]: │ in _pre_forward                                                              │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    378 │   │   if handle:                                                    │
[rank7]: [rank7]: [rank7]: │    379 │   │   │   handle._training_state = HandleTrainingState.FORWARD      │
[rank7]: [rank7]: [rank7]: │    380 │   │   if unshard_fn is not None:                                    │
[rank7]: [rank7]: [rank7]: │ ❱  381 │   │   │   unshard_fn(state, handle)                                 │
[rank7]: [rank7]: [rank7]: │    382 │   │   # Register post-backward hooks to reshard the parameters and  │
[rank7]: [rank7]: [rank7]: │    383 │   │   # their gradients. They must be re-registered every forward p │
[rank7]: [rank7]: [rank7]: │    384 │   │   # the `grad_fn` is mutated.                                   │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/distributed/fsdp/_runtime_utils.py:416  │
[rank7]: [rank7]: [rank7]: │ in _pre_forward_unshard                                                      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    413 │   # If the handles have been prefetched, then there is no need to c │
[rank7]: [rank7]: [rank7]: │    414 │   # `_unshard()` again                                              │
[rank7]: [rank7]: [rank7]: │    415 │   if not handle._prefetched:                                        │
[rank7]: [rank7]: [rank7]: │ ❱  416 │   │   _unshard(state, handle, state._unshard_stream, state._pre_uns │
[rank7]: [rank7]: [rank7]: │    417 │   handle._needs_pre_forward_unshard = False                         │
[rank7]: [rank7]: [rank7]: │    418 │   # Don't wait during trace                                         │
[rank7]: [rank7]: [rank7]: │    419 │   if not torch.distributed._functional_collectives.is_torchdynamo_c │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/distributed/fsdp/_runtime_utils.py:300  │
[rank7]: [rank7]: [rank7]: │ in _unshard                                                                  │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    297 │   │   │   ):                                                        │
[rank7]: [rank7]: [rank7]: │    298 │   │   │   │   event.synchronize()                                   │
[rank7]: [rank7]: [rank7]: │    299 │   with state._device_handle.stream(unshard_stream):                 │
[rank7]: [rank7]: [rank7]: │ ❱  300 │   │   handle.unshard()                                              │
[rank7]: [rank7]: [rank7]: │    301 │   │   handle.post_unshard()                                         │
[rank7]: [rank7]: [rank7]: │    302                                                                       │
[rank7]: [rank7]: [rank7]: │    303                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/trainer/_patch_pytorch.py:1009 in    │
[rank7]: [rank7]: [rank7]: │ unshard_with_sync                                                            │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   1006 │   from torch.distributed.fsdp._flat_param import FlatParamHandle    │
[rank7]: [rank7]: [rank7]: │   1007 │   original_unshard = FlatParamHandle.unshard                        │
[rank7]: [rank7]: [rank7]: │   1008 │                                                                     │
[rank7]: [rank7]: [rank7]: │ ❱ 1009 │   @no_type_check                                                    │
[rank7]: [rank7]: [rank7]: │   1010 │   def unshard_with_sync(self):                                      │
[rank7]: [rank7]: [rank7]: │   1011 │   │   """Run the unshard logic, but with a sync after a :meth:`_all │
[rank7]: [rank7]: [rank7]: │   1012                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/composer/trainer/_patch_pytorch.py:1039 in    │
[rank7]: [rank7]: [rank7]: │ torch_dynamo_resume_in_unshard_with_sync_at_1039                             │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   1036 │   │   found_cuda_oom_tensor = torch.tensor([0], dtype=torch.uint8). │
[rank7]: [rank7]: [rank7]: │   1037 │   │                                                                 │
[rank7]: [rank7]: [rank7]: │   1038 │   │   dist.all_reduce(found_cuda_oom_tensor, reduce_operation='MAX' │
[rank7]: [rank7]: [rank7]: │ ❱ 1039 │   │   found_cuda_oom = found_cuda_oom_tensor.item()                 │
[rank7]: [rank7]: [rank7]: │   1040 │   │   # Signal current rank is still in batch                       │
[rank7]: [rank7]: [rank7]: │   1041 │   │   all_ranks_finished_tensor = torch.tensor([0], dtype=torch.uin │
[rank7]: [rank7]: [rank7]: │   1042                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_dynamo/eval_frame.py:600 in _fn        │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    597 │   │   def _fn(*args, **kwargs):                                     │
[rank7]: [rank7]: [rank7]: │    598 │   │   │   prior = set_eval_frame(callback)                          │
[rank7]: [rank7]: [rank7]: │    599 │   │   │   try:                                                      │
[rank7]: [rank7]: [rank7]: │ ❱  600 │   │   │   │   return fn(*args, **kwargs)                            │
[rank7]: [rank7]: [rank7]: │    601 │   │   │   finally:                                                  │
[rank7]: [rank7]: [rank7]: │    602 │   │   │   │   set_eval_frame(prior)                                 │
[rank7]: [rank7]: [rank7]: │    603                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py:987 in       │
[rank7]: [rank7]: [rank7]: │ forward                                                                      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    984 │   │   full_args = []                                                │
[rank7]: [rank7]: [rank7]: │    985 │   │   full_args.extend(params_flat)                                 │
[rank7]: [rank7]: [rank7]: │    986 │   │   full_args.extend(runtime_args)                                │
[rank7]: [rank7]: [rank7]: │ ❱  987 │   │   return compiled_fn(full_args)                                 │
[rank7]: [rank7]: [rank7]: │    988 │                                                                     │
[rank7]: [rank7]: [rank7]: │    989 │   # Just for convenience                                            │
[rank7]: [rank7]: [rank7]: │    990 │   forward.zero_grad = mod.zero_grad                                 │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappe │
[rank7]: [rank7]: [rank7]: │ rs.py:217 in runtime_wrapper                                                 │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    214 │   │   │   try:                                                      │
[rank7]: [rank7]: [rank7]: │    215 │   │   │   │   if grad_enabled:                                      │
[rank7]: [rank7]: [rank7]: │    216 │   │   │   │   │   torch._C._set_grad_enabled(False)                 │
[rank7]: [rank7]: [rank7]: │ ❱  217 │   │   │   │   all_outs = call_func_at_runtime_with_args(            │
[rank7]: [rank7]: [rank7]: │    218 │   │   │   │   │   compiled_fn, args, disable_amp=disable_amp, steal │
[rank7]: [rank7]: [rank7]: │    219 │   │   │   │   )                                                     │
[rank7]: [rank7]: [rank7]: │    220 │   │   │   finally:                                                  │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_functorch/_aot_autograd/utils.py:120   │
[rank7]: [rank7]: [rank7]: │ in call_func_at_runtime_with_args                                            │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   117 │   context = torch._C._DisableAutocast if disable_amp else nullcontex │
[rank7]: [rank7]: [rank7]: │   118 │   with context():                                                    │
[rank7]: [rank7]: [rank7]: │   119 │   │   if hasattr(f, "_boxed_call"):                                  │
[rank7]: [rank7]: [rank7]: │ ❱ 120 │   │   │   out = normalize_as_list(f(args))                           │
[rank7]: [rank7]: [rank7]: │   121 │   │   else:                                                          │
[rank7]: [rank7]: [rank7]: │   122 │   │   │   # TODO: Please remove soon                                 │
[rank7]: [rank7]: [rank7]: │   123 │   │   │   # https://github.com/pytorch/pytorch/pull/83137#issuecomme │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappe │
[rank7]: [rank7]: [rank7]: │ rs.py:451 in wrapper                                                         │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    448 │   │   │   │   │   runtime_metadata.num_forward_returns,             │
[rank7]: [rank7]: [rank7]: │    449 │   │   │   │   )                                                     │
[rank7]: [rank7]: [rank7]: │    450 │   │   │   │   return out                                            │
[rank7]: [rank7]: [rank7]: │ ❱  451 │   │   │   return compiled_fn(runtime_args)                          │
[rank7]: [rank7]: [rank7]: │    452 │   │                                                                 │
[rank7]: [rank7]: [rank7]: │    453 │   │   return wrapper                                                │
[rank7]: [rank7]: [rank7]: │    454                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/codecache.py:1131 in __call__ │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   1128 │                                                                     │
[rank7]: [rank7]: [rank7]: │   1129 │   def __call__(self, inputs: List[Any]) -> Any:                     │
[rank7]: [rank7]: [rank7]: │   1130 │   │   assert self.current_callable is not None                      │
[rank7]: [rank7]: [rank7]: │ ❱ 1131 │   │   return self.current_callable(inputs)                          │
[rank7]: [rank7]: [rank7]: │   1132                                                                       │
[rank7]: [rank7]: [rank7]: │   1133                                                                       │
[rank7]: [rank7]: [rank7]: │   1134 def cpp_compiler() -> str:                                            │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/compile_fx.py:944 in run      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    941 │                                                                     │
[rank7]: [rank7]: [rank7]: │    942 │   def run(new_inputs):                                              │
[rank7]: [rank7]: [rank7]: │    943 │   │   copy_misaligned_inputs(new_inputs, inputs_to_check)           │
[rank7]: [rank7]: [rank7]: │ ❱  944 │   │   return model(new_inputs)                                      │
[rank7]: [rank7]: [rank7]: │    945 │                                                                     │
[rank7]: [rank7]: [rank7]: │    946 │   return run                                                        │
[rank7]: [rank7]: [rank7]: │    947                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /tmp/torchinductor_root/p7/cp7mtlyomhhf7t4rpjcvw6y456kkfpckmd5xryemyrunauchz │
[rank7]: [rank7]: [rank7]: │ bcv.py:84 in call                                                            │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    81 │   │   del arg0_1                                                     │
[rank7]: [rank7]: [rank7]: │    82 │   │   # Source Nodes: [], Original ATen: []                          │
[rank7]: [rank7]: [rank7]: │    83 │   │   stream7 = get_raw_stream(7)                                    │
[rank7]: [rank7]: [rank7]: │ ❱  84 │   │   triton_poi_fused_0.run(buf0, arg1_1, 212553216, grid=grid(2125 │
[rank7]: [rank7]: [rank7]: │    85 │   │   del arg1_1                                                     │
[rank7]: [rank7]: [rank7]: │    86 │   │   del buf0                                                       │
[rank7]: [rank7]: [rank7]: │    87 │   return ()                                                          │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/runtime/triton_heuristics.py: │
[rank7]: [rank7]: [rank7]: │ 820 in run                                                                   │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    817 │   │   │   if len(self.launchers) == 0:                              │
[rank7]: [rank7]: [rank7]: │    818 │   │   │   │   self.precompile()                                     │
[rank7]: [rank7]: [rank7]: │    819 │   │   │   if len(self.launchers) > 1:                               │
[rank7]: [rank7]: [rank7]: │ ❱  820 │   │   │   │   self.autotune_to_one_config(*args, grid=grid, **kwarg │
[rank7]: [rank7]: [rank7]: │    821 │   │                                                                 │
[rank7]: [rank7]: [rank7]: │    822 │   │   if not getattr(                                               │
[rank7]: [rank7]: [rank7]: │    823 │   │   │   self.launchers[0].config, "found_by_coordesc", False      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/runtime/triton_heuristics.py: │
[rank7]: [rank7]: [rank7]: │ 718 in autotune_to_one_config                                                │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    715 │   def autotune_to_one_config(self, *args, **kwargs):                │
[rank7]: [rank7]: [rank7]: │    716 │   │   """Do the actual autotuning"""                                │
[rank7]: [rank7]: [rank7]: │    717 │   │   start_time = time.time_ns()                                   │
[rank7]: [rank7]: [rank7]: │ ❱  718 │   │   timings = self.benchmark_all_configs(*args, **kwargs)         │
[rank7]: [rank7]: [rank7]: │    719 │   │   time_taken_ns = time.time_ns() - start_time                   │
[rank7]: [rank7]: [rank7]: │    720 │   │   self.launchers = [builtins.min(timings, key=timings.get)]     │
[rank7]: [rank7]: [rank7]: │    721 │   │   if self.save_cache_hook:                                      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_dynamo/utils.py:231 in time_wrapper    │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    228 │   │   │   try:                                                      │
[rank7]: [rank7]: [rank7]: │    229 │   │   │   │   with torch.profiler.record_function(f"{key} (dynamo_t │
[rank7]: [rank7]: [rank7]: │    230 │   │   │   │   │   t0 = time.time()                                  │
[rank7]: [rank7]: [rank7]: │ ❱  231 │   │   │   │   │   r = func(*args, **kwargs)                         │
[rank7]: [rank7]: [rank7]: │    232 │   │   │   │   │   time_spent = time.time() - t0                     │
[rank7]: [rank7]: [rank7]: │    233 │   │   │   │   compilation_time_metrics[key].append(time_spent)      │
[rank7]: [rank7]: [rank7]: │    234 │   │   │   except Exception as e:                                    │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/runtime/triton_heuristics.py: │
[rank7]: [rank7]: [rank7]: │ 693 in benchmark_all_configs                                                 │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    690 │                                                                     │
[rank7]: [rank7]: [rank7]: │    691 │   @dynamo_timed                                                     │
[rank7]: [rank7]: [rank7]: │    692 │   def benchmark_all_configs(self, *args, **kwargs):                 │
[rank7]: [rank7]: [rank7]: │ ❱  693 │   │   timings = {                                                   │
[rank7]: [rank7]: [rank7]: │    694 │   │   │   launcher: self.bench(launcher, *args, **kwargs)           │
[rank7]: [rank7]: [rank7]: │    695 │   │   │   for launcher in self.launchers                            │
[rank7]: [rank7]: [rank7]: │    696 │   │   }                                                             │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/runtime/triton_heuristics.py: │
[rank7]: [rank7]: [rank7]: │ 694 in <dictcomp>                                                            │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    691 │   @dynamo_timed                                                     │
[rank7]: [rank7]: [rank7]: │    692 │   def benchmark_all_configs(self, *args, **kwargs):                 │
[rank7]: [rank7]: [rank7]: │    693 │   │   timings = {                                                   │
[rank7]: [rank7]: [rank7]: │ ❱  694 │   │   │   launcher: self.bench(launcher, *args, **kwargs)           │
[rank7]: [rank7]: [rank7]: │    695 │   │   │   for launcher in self.launchers                            │
[rank7]: [rank7]: [rank7]: │    696 │   │   }                                                             │
[rank7]: [rank7]: [rank7]: │    697                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/runtime/triton_heuristics.py: │
[rank7]: [rank7]: [rank7]: │ 665 in bench                                                                 │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    662 │   │   │   │   stream=stream,                                        │
[rank7]: [rank7]: [rank7]: │    663 │   │   │   )                                                         │
[rank7]: [rank7]: [rank7]: │    664 │   │                                                                 │
[rank7]: [rank7]: [rank7]: │ ❱  665 │   │   return do_bench_gpu(kernel_call, rep=40, fast_flush=True)     │
[rank7]: [rank7]: [rank7]: │    666 │                                                                     │
[rank7]: [rank7]: [rank7]: │    667 │   def clone_args(self, *args, **kwargs) -> Tuple[List[Any], Dict[st │
[rank7]: [rank7]: [rank7]: │    668 │   │   from ..compile_fx import clone_preserve_strides               │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/runtime/runtime_utils.py:113  │
[rank7]: [rank7]: [rank7]: │ in do_bench_gpu                                                              │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   110 │                                                                      │
[rank7]: [rank7]: [rank7]: │   111 │   if quantile_field_name not in kwargs:                              │
[rank7]: [rank7]: [rank7]: │   112 │   │   kwargs[quantile_field_name] = (0.5, 0.2, 0.8)                  │
[rank7]: [rank7]: [rank7]: │ ❱ 113 │   return triton_do_bench(*args, **kwargs)[0]                         │
[rank7]: [rank7]: [rank7]: │   114                                                                        │
[rank7]: [rank7]: [rank7]: │   115                                                                        │
[rank7]: [rank7]: [rank7]: │   116 def do_bench_cpu(fn, warmup=5, times=20):                              │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/triton/testing.py:103 in do_bench             │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │   100 │   assert return_mode in ["min", "max", "mean", "median"]             │
[rank7]: [rank7]: [rank7]: │   101 │   import torch                                                       │
[rank7]: [rank7]: [rank7]: │   102 │                                                                      │
[rank7]: [rank7]: [rank7]: │ ❱ 103 │   fn()                                                               │
[rank7]: [rank7]: [rank7]: │   104 │   torch.cuda.synchronize()                                           │
[rank7]: [rank7]: [rank7]: │   105 │                                                                      │
[rank7]: [rank7]: [rank7]: │   106 │   # We maintain a buffer of 256 MB that we clear                     │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/runtime/triton_heuristics.py: │
[rank7]: [rank7]: [rank7]: │ 657 in kernel_call                                                           │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    654 │   │   │   │   │   {**dict(zip(self.arg_names, args)), **launcher.co │
[rank7]: [rank7]: [rank7]: │    655 │   │   │   │   )                                                     │
[rank7]: [rank7]: [rank7]: │    656 │   │   │                                                             │
[rank7]: [rank7]: [rank7]: │ ❱  657 │   │   │   cloned_args, cloned_kwargs = self.clone_args(*args, **kwa │
[rank7]: [rank7]: [rank7]: │    658 │   │   │   launcher(                                                 │
[rank7]: [rank7]: [rank7]: │    659 │   │   │   │   *cloned_args,                                         │
[rank7]: [rank7]: [rank7]: │    660 │   │   │   │   **cloned_kwargs,                                      │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/runtime/triton_heuristics.py: │
[rank7]: [rank7]: [rank7]: │ 677 in clone_args                                                            │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    674 │   │   for i, arg in enumerate(args):                                │
[rank7]: [rank7]: [rank7]: │    675 │   │   │   if self.fn.arg_names[i] in self.mutated_arg_names:        │
[rank7]: [rank7]: [rank7]: │    676 │   │   │   │   assert isinstance(arg, torch.Tensor)                  │
[rank7]: [rank7]: [rank7]: │ ❱  677 │   │   │   │   cloned_args.append(clone_preserve_strides(arg))       │
[rank7]: [rank7]: [rank7]: │    678 │   │   │   else:                                                     │
[rank7]: [rank7]: [rank7]: │    679 │   │   │   │   cloned_args.append(arg)                               │
[rank7]: [rank7]: [rank7]: │    680                                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │ /usr/lib/python3/dist-packages/torch/_inductor/compile_fx.py:889 in          │
[rank7]: [rank7]: [rank7]: │ clone_preserve_strides                                                       │
[rank7]: [rank7]: [rank7]: │                                                                              │
[rank7]: [rank7]: [rank7]: │    886 │   needed_size = (                                                   │
[rank7]: [rank7]: [rank7]: │    887 │   │   sum((shape - 1) * stride for shape, stride in zip(x.size(), x │
[rank7]: [rank7]: [rank7]: │    888 │   )                                                                 │
[rank7]: [rank7]: [rank7]: │ ❱  889 │   buffer = torch.as_strided(x, (needed_size,), (1,)).clone()        │
[rank7]: [rank7]: [rank7]: │    890 │   return torch.as_strided(buffer, x.size(), x.stride())             │
[rank7]: [rank7]: [rank7]: │    891                                                                       │
[rank7]: [rank7]: [rank7]: │    892                                                                       │
[rank7]: [rank7]: [rank7]: ╰──────────────────────────────────────────────────────────────────────────────╯
[rank7]: [rank7]: [rank7]: RuntimeError: setStorage: sizes [212553216], strides [1], storage offset 0, and
[rank7]: [rank7]: [rank7]: itemsize 2 requiring a storage size of 425106432 are out of bounds for storage
[rank7]: [rank7]: [rank7]: of size 0
[rank7]:[W1029 15:44:28.411663930 Functional.cpp:45] Warning: At the time of process termination, there are still 4 unwaited c10d_functional collective calls. Please review your program to ensure c10d_functional.wait_tensor() is invoked on all tensors returned from c10d_functional collective ops before they are used. (function ~WorkRegistry)
```
### Minified repro
Architecture is simple MosaicBERT architecture
_No response_

### Versions

2.4.1 pip version ran through composer

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @zhaojuanmao @mrshenli @rohan-varma @chauhang @ezyang @penguinwu