BF16 stochastic rounding does not work distributed (FSDP)

[nathan-az]

Error arose using torchtune. Config and stacktrace are below. Testing was done on a feature branch but with no non-standard features enabled. It works fine with nproc=1, but with FSDP or TP enabled, it breaks.

torchtune config

batch_size: 1
checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_files:
    filename_format: model-{}-of-{}.safetensors
    max_filename: '00004'
  model_type: LLAMA3
  output_dir: ${output_dir}
  recipe_checkpoint: null
  checkpoint_dir: models/llama_3_1_8b
clip_grad_norm: null
compile: false
custom_sharded_layers: []
data_parallel_shard_dim: 2
data_parallel_replicate_dim: 1
tensor_parallel_dim: 1
enable_fp8_training: false
enable_loss_parallel: true
tensor_parallel_plan:
  _component_: torchtune.models.llama3.base_llama_tp_plan
dataset:
  _component_: torchtune.datasets.alpaca_dataset
  packed: true
  train_on_input: true
  split: train[:10%]
device: cuda
dtype: bf16
enable_activation_checkpointing: true
enable_activation_offloading: false
activation_offloading_use_streams: false
epochs: 1
fsdp_cpu_offload: false
gradient_accumulation_steps: 1
log_every_n_steps: 1
log_level: INFO
log_peak_memory_stats: true
loss:
  _component_: torchtune.modules.loss.LinearCrossEntropyLoss
max_steps_per_epoch: 10
metric_logger:
  _component_: torchtune.training.metric_logging.MLFlowLogger
  experiment_name: llama_3_1_8b_debug
  run_name: debug_stochastic_round_dp1
optimizer:
  _component_: torchao.optim.AdamW8bit
  bf16_stochastic_round: true
  lr: 4.0e-05
optimizer_in_bwd: false
output_dir: outputs
resume_from_checkpoint: false
seed: 100
shuffle: true
tokenizer:
  max_seq_len: 8192
  path: models/llama_3_1_8b/original/tokenizer.model
  _component_: torchtune.models.llama3.llama3_tokenizer
model:
  _component_: torchtune.models.llama3_1.llama3_1_8b

Stacktrace

[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/torchtune/recipes/full_finetune_distributed.py", line 1118, in <module>
[rank0]:     sys.exit(recipe_main())
[rank0]:              ^^^^^^^^^^^^^
[rank0]:   File "/workspace/torchtune/torchtune/config/_parse.py", line 99, in wrapper
[rank0]:     sys.exit(recipe_main(conf))
[rank0]:              ^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/torchtune/recipes/full_finetune_distributed.py", line 1113, in recipe_main
[rank0]:     recipe.train()
[rank0]:   File "/workspace/torchtune/recipes/full_finetune_distributed.py", line 982, in train
[rank0]:     self._optimizer.step()
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/optim/optimizer.py", line 506, in wrapper
[rank0]:     out = func(*args, **kwargs)
[rank0]:           ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torchao/optim/adam.py", line 134, in step
[rank0]:     torch.compile(single_param_adam, fullgraph=True, dynamic=False)(
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 699, in compile_wrapper
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 1469, in __call__
[rank0]:     return self._torchdynamo_orig_callable(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 625, in __call__
[rank0]:     return _compile(
[rank0]:            ^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 1092, in _compile
[rank0]:     guarded_code = compile_inner(code, one_graph, hooks, transform)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_utils_internal.py", line 97, in wrapper_function
[rank0]:     return function(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 779, in compile_inner
[rank0]:     return _compile_inner(code, one_graph, hooks, transform)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 818, in _compile_inner
[rank0]:     out_code = transform_code_object(code, transform)
[rank0]:                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1424, in transform_code_object
[rank0]:     transformations(instructions, code_options)
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 265, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 743, in transform
[rank0]:     tracer.run()
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 3484, in run
[rank0]:     super().run()
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1359, in run
[rank0]:     while self.step():
[rank0]:           ^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1263, in step
[rank0]:     self.dispatch_table[inst.opcode](self, inst)
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 831, in wrapper
[rank0]:     return inner_fn(self, inst)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2903, in CALL
[rank0]:     self._call(inst)
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2897, in _call
[rank0]:     self.call_function(fn, args, kwargs)
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1189, in call_function
[rank0]:     self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/lazy.py", line 201, in realize_and_forward
[rank0]:     return getattr(self.realize(), name)(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 516, in call_function
[rank0]:     return super().call_function(tx, args, kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 291, in call_function
[rank0]:     return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1206, in inline_user_function_return
[rank0]:     return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 3678, in inline_call
[rank0]:     return tracer.inline_call_()
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 3881, in inline_call_
[rank0]:     self.run()
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1359, in run
[rank0]:     while self.step():
[rank0]:           ^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1263, in step
[rank0]:     self.dispatch_table[inst.opcode](self, inst)
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2161, in COMPARE_OP
[rank0]:     self.push(compare_op_handlers[inst.argval](self, self.popn(2), {}))
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/builtin.py", line 1168, in call_function
[rank0]:     return handler(tx, args, kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/builtin.py", line 1128, in _handle_insert_op_in_graph
[rank0]:     return wrap_fx_proxy(tx, proxy)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 2500, in wrap_fx_proxy
[rank0]:     return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 2566, in wrap_fx_proxy_cls
[rank0]:     return _wrap_fx_proxy(
[rank0]:            ^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 2664, in _wrap_fx_proxy
[rank0]:     example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3347, in get_fake_value
[rank0]:     raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3245, in get_fake_value
[rank0]:     ret_val = wrap_fake_exception(
[rank0]:               ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 2745, in wrap_fake_exception
[rank0]:     return fn()
[rank0]:            ^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3246, in <lambda>
[rank0]:     lambda: run_node(tx.output, node, args, kwargs, nnmodule)
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3454, in run_node
[rank0]:     raise RuntimeError(make_error_message(e)).with_traceback(
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3413, in run_node
[rank0]:     return node.target(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_compile.py", line 51, in inner
[rank0]:     return disable_fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 893, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/distributed/tensor/_api.py", line 350, in __torch_dispatch__
[rank0]:     return DTensor._op_dispatcher.dispatch(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/distributed/tensor/_dispatch.py", line 157, in dispatch
[rank0]:     op_info = self.unwrap_to_op_info(op_call, args, kwargs)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/distributed/tensor/_dispatch.py", line 356, in unwrap_to_op_info
[rank0]:     self._try_replicate_spec_for_scalar_tensor(
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torch/distributed/tensor/_dispatch.py", line 458, in _try_replicate_spec_for_scalar_tensor
[rank0]:     raise RuntimeError(
[rank0]: torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function <built-in function lt>(*(FakeTensor(..., device='cuda:0', size=(128256, 4096), dtype=torch.int32), DTensor(local_tensor=FakeTensor(..., device='cuda:0', size=(64128, 4096), dtype=torch.int32), device_mesh=DeviceMesh('cuda', [0, 1], mesh_dim_names=('dp_shard',)), placements=(Shard(dim=0),))), **{}): got RuntimeError('aten.lt.Tensor: got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!')

[rank0]: from user code:
[rank0]:    File "/opt/conda/lib/python3.11/site-packages/torchao/optim/adam.py", line 199, in single_param_adam
[rank0]:     p.copy_(_fp32_to_bf16_sr(p_f32))
[rank0]:   File "/opt/conda/lib/python3.11/site-packages/torchao/optim/quant_utils.py", line 138, in _fp32_to_bf16_sr
[rank0]:     rand_16bit < x_fraction,  # this is True with the probability of p_fraction

[gau-nernst]

@nathan-az what's ur pytorch version?

[nathan-az]

@gau-nernst latest (or a day old maximum):

torch==2.8.0.dev20250603+cu128
torchao==0.12.0.dev20250603+cu128

[gau-nernst]

@nathan-az Can you try changing this line

ao/torchao/optim/quant_utils.py

Lines 130 to 132 in 3aa9361

	rand_16bit = torch.randint(
	0, 1 << 16, x_f32.shape, device=x_f32.device, dtype=torch.int32
	)

To

rand_16bit = torch.randint_like(x_f32, 0, 1 << 16, dtype=torch.int32)

[nathan-az]

No luck @gau-nernst :

[rank1]: Traceback (most recent call last):
[rank1]:   File "/workspace/torchtune/recipes/full_finetune_distributed.py", line 1118, in <module>
[rank1]:     sys.exit(recipe_main())
[rank1]:              ^^^^^^^^^^^^^
[rank1]:   File "/workspace/torchtune/torchtune/config/_parse.py", line 99, in wrapper
[rank1]:     sys.exit(recipe_main(conf))
[rank1]:              ^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/torchtune/recipes/full_finetune_distributed.py", line 1113, in recipe_main
[rank1]:     recipe.train()
[rank1]:   File "/workspace/torchtune/recipes/full_finetune_distributed.py", line 982, in train
[rank1]:     self._optimizer.step()
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/optim/optimizer.py", line 506, in wrapper
[rank1]:     out = func(*args, **kwargs)
[rank1]:           ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank1]:     return func(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torchao/optim/adam.py", line 134, in step
[rank1]:     torch.compile(single_param_adam, fullgraph=True, dynamic=False)(
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 699, in compile_wrapper
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 1469, in __call__
[rank1]:     return self._torchdynamo_orig_callable(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 625, in __call__
[rank1]:     return _compile(
[rank1]:            ^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 1092, in _compile
[rank1]:     guarded_code = compile_inner(code, one_graph, hooks, transform)
[rank1]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_utils_internal.py", line 97, in wrapper_function
[rank1]:     return function(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 779, in compile_inner
[rank1]:     return _compile_inner(code, one_graph, hooks, transform)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 818, in _compile_inner
[rank1]:     out_code = transform_code_object(code, transform)
[rank1]:                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1424, in transform_code_object
[rank1]:     transformations(instructions, code_options)
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 265, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 743, in transform
[rank1]:     tracer.run()
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 3484, in run
[rank1]:     super().run()
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1359, in run
[rank1]:     while self.step():
[rank1]:           ^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1263, in step
[rank1]:     self.dispatch_table[inst.opcode](self, inst)
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 831, in wrapper
[rank1]:     return inner_fn(self, inst)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2903, in CALL
[rank1]:     self._call(inst)
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2897, in _call
[rank1]:     self.call_function(fn, args, kwargs)
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1189, in call_function
[rank1]:     self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/lazy.py", line 201, in realize_and_forward
[rank1]:     return getattr(self.realize(), name)(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 516, in call_function
[rank1]:     return super().call_function(tx, args, kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 291, in call_function
[rank1]:     return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1206, in inline_user_function_return
[rank1]:     return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 3678, in inline_call
[rank1]:     return tracer.inline_call_()
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 3881, in inline_call_
[rank1]:     self.run()
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1359, in run
[rank1]:     while self.step():
[rank1]:           ^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1263, in step
[rank1]:     self.dispatch_table[inst.opcode](self, inst)
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 831, in wrapper
[rank1]:     return inner_fn(self, inst)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2903, in CALL
[rank1]:     self._call(inst)
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2897, in _call
[rank1]:     self.call_function(fn, args, kwargs)
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1189, in call_function
[rank1]:     self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/lazy.py", line 201, in realize_and_forward
[rank1]:     return getattr(self.realize(), name)(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/torch.py", line 1297, in call_function
[rank1]:     tensor_variable = wrap_fx_proxy(
[rank1]:                       ^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 2500, in wrap_fx_proxy
[rank1]:     return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 2566, in wrap_fx_proxy_cls
[rank1]:     return _wrap_fx_proxy(
[rank1]:            ^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 2664, in _wrap_fx_proxy
[rank1]:     example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
[rank1]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3347, in get_fake_value
[rank1]:     raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3245, in get_fake_value
[rank1]:     ret_val = wrap_fake_exception(
[rank1]:               ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 2745, in wrap_fake_exception
[rank1]:     return fn()
[rank1]:            ^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3246, in <lambda>
[rank1]:     lambda: run_node(tx.output, node, args, kwargs, nnmodule)
[rank1]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3454, in run_node
[rank1]:     raise RuntimeError(make_error_message(e)).with_traceback(
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 3413, in run_node
[rank1]:     return node.target(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_compile.py", line 51, in inner
[rank1]:     return disable_fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 893, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/distributed/tensor/_api.py", line 350, in __torch_dispatch__
[rank1]:     return DTensor._op_dispatcher.dispatch(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/distributed/tensor/_dispatch.py", line 190, in dispatch
[rank1]:     random._rng_tracker = random.OffsetBasedRNGTracker(mesh)
[rank1]:                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/distributed/tensor/_random.py", line 179, in __init__
[rank1]:     rng_state = self._device_handle.get_rng_state().to(self._device)
[rank1]:                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/utils/_stats.py", line 27, in wrapper
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1346, in __torch_dispatch__
[rank1]:     return self.dispatch(func, types, args, kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 2034, in dispatch
[rank1]:     return self._cached_dispatch_impl(func, types, args, kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1442, in _cached_dispatch_impl
[rank1]:     return self._dispatch_impl(func, types, args, kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 2328, in _dispatch_impl
[rank1]:     (flat_args, flat_arg_fake_tensors) = self.validate_and_convert_non_fake_tensors(
[rank1]:                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 2779, in validate_and_convert_non_fake_tensors
[rank1]:     validated_args = [validate(a) for a in flat_args]
[rank1]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 2779, in <listcomp>
[rank1]:     validated_args = [validate(a) for a in flat_args]
[rank1]:                       ^^^^^^^^^^^
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 2767, in validate
[rank1]:     raise AssertionError(
[rank1]: torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function <built-in method randint_like of type object at 0x72c22ceca340>(*(DTensor(local_tensor=FakeTensor(..., device='cuda:1', size=(64128, 4096)), device_mesh=DeviceMesh('cuda', [0, 1], mesh_dim_names=('dp_shard',)), placements=(Shard(dim=0),)), 0, 65536), **{'dtype': torch.int32}): got AssertionError("Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.to.dtype_layout(tensor([...], size=(16,), dtype=torch.uint8), dtype=torch.uint8, layout=torch.strided, device=device(type='cuda', index=1))")

[rank1]: from user code:
[rank1]:    File "/opt/conda/lib/python3.11/site-packages/torchao/optim/adam.py", line 199, in single_param_adam
[rank1]:     p.copy_(_fp32_to_bf16_sr(p_f32))
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torchao/optim/quant_utils.py", line 133, in _fp32_to_bf16_sr
[rank1]:     rand_16bit = torch.randint_like(x_f32, 0, 1 << 16, dtype=torch.int32)

Key part appears to be:

[rank1]: torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function <built-in method randint_like of type object at 0x72c22ceca340>(*(DTensor(local_tensor=FakeTensor(..., device='cuda:1', size=(64128, 4096)), device_mesh=DeviceMesh('cuda', [0, 1], mesh_dim_names=('dp_shard',)), placements=(Shard(dim=0),)), 0, 65536), **{'dtype': torch.int32}): got AssertionError("Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.to.dtype_layout(tensor([...], size=(16,), dtype=torch.uint8), dtype=torch.uint8, layout=torch.strided, device=device(type='cuda', index=1))")

[rank1]: from user code:
[rank1]:    File "/opt/conda/lib/python3.11/site-packages/torchao/optim/adam.py", line 199, in single_param_adam
[rank1]:     p.copy_(_fp32_to_bf16_sr(p_f32))
[rank1]:   File "/opt/conda/lib/python3.11/site-packages/torchao/optim/quant_utils.py", line 133, in _fp32_to_bf16_sr
[rank1]:     rand_16bit = torch.randint_like(x_f32, 0, 1 << 16, dtype=torch.int32)

Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.to.dtype_layout(tensor([...], size=(16,), dtype=torch.uint8), dtype=torch.uint8, layout=torch.strided, device=device(type='cuda', index=1))"

Not sure how to do this (or if this is an actual valid solution)

[gau-nernst]

I was hoping the fix would be simple as that. Will investigate more when I have the chance.

But it does look like a torch.compile+DTensor bug i.e. torch.randint_like() doesn't work with DTensor+torch.compile. cc @bdhirsh

I likely can come up with some "not so nice" hack (e.g. check for DTensor explicitly). @nathan-az Just curious, have you tried / do u encounter the error with torch2.7?

[nathan-az]

I haven't tried 2.7.0 - I'm using some other experimental features concurrently which prohibit moving off nightly unfortunately.

I tried constructing the DTensor explicitly (probably same end result as yours) but it threw a similar error, so I'm at a loss unfortunately.

[bdhirsh]

@nathan-az any chance you have a smaller repro I can run that involves constructing the DTensor explicitly and raises the error you mentioned?

[nathan-az]

Sure @bdhirsh . This one doesn't actually use the optimiser, but shows that the _fp32_to_bf16_sr function doesn't work with DTensor inputs and raises the same error.

import os
import torch
import torch.distributed as dist
from torch.distributed.tensor import Shard, DTensor
from torch.distributed.device_mesh import init_device_mesh
from torchao.optim.quant_utils import _fp32_to_bf16_sr


def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)


def cleanup():
    dist.destroy_process_group()


def main():
    """Main training function."""
    # Get distributed training info from environment variables set by torchrun
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    
    setup(local_rank, world_size)
    
    mesh = init_device_mesh("cuda", (world_size,))
    x = torch.randn((64,)).cuda()
    distx = DTensor.from_local(x, device_mesh=mesh, placements=[Shard(0)])
    rounded = _fp32_to_bf16_sr(distx)
    cleanup()


if __name__ == "__main__":
    main()

[bdhirsh]

Hmm, telling DTensor to allow "implicit replication" gets your repro to pass for me. Any chance it's possible to try it on the larger FSDP repro?

from torch.distributed.tensor.experimental import implicit_replication

with implicit_replication()
    distx = DTensor.from_local(x, device_mesh=mesh, placements=[Shard(0)])
    rounded = _fp32_to_bf16_sr(distx)

Side note - if this works, it is a bit annoying that the error message didn't make it clear that this is the way to fix it, hopefully we can make that better

[gau-nernst]

@bdhirsh I haven't had the time to debug it further, but is using torch.randint_like() not a valid solution? Though @nathan-az is getting an error with Dynamo #2296 (comment) with it

[nathan-az]

Hmm I can look further in the (Australian) morning @bdhirsh - however I'm not sure if this is a desirable solution. If we consider the large model and tensor parallel case, AFAIK the optimiser states should be safe to remain sharded across workers.

Wouldn't implicit replication require an all-gather? This could be relatively expensive in both time and memory usage (if my understanding is correct that optimiser states shard across the same placements as their respective parameters, and don't currently require additional collectives when stepping).

[bdhirsh]

oh yep, thanks - I get the same error when I make that change locally (stacktrace below).

I should have remembered, but this is actually a known problem that DTensor's specific implementation of RNG tracking is not compile friendly, which means that you can't use RNG ops with DTensor and compile. This seems like good data for prioritizing though (cc @XilunWu / @wconstab ).

stacktrace:

[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/_dynamo/eval_frame.py", line 893, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/distributed/tensor/_api.py", line 350, in __torch_dispatch__
[rank0]:     return DTensor._op_dispatcher.dispatch(
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/distributed/tensor/_dispatch.py", line 190, in dispatch
[rank0]:     random._rng_tracker = random.OffsetBasedRNGTracker(mesh)
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/distributed/tensor/_random.py", line 179, in __init__
[rank0]:     rng_state = self._device_handle.get_rng_state().to(self._device)
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/utils/_stats.py", line 27, in wrapper
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/_subclasses/fake_tensor.py", line 1346, in __torch_dispatch__
[rank0]:     return self.dispatch(func, types, args, kwargs)
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/_subclasses/fake_tensor.py", line 2034, in dispatch
[rank0]:     return self._cached_dispatch_impl(func, types, args, kwargs)
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/_subclasses/fake_tensor.py", line 1442, in _cached_dispatch_impl
[rank0]:     return self._dispatch_impl(func, types, args, kwargs)
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/_subclasses/fake_tensor.py", line 2328, in _dispatch_impl
[rank0]:     (flat_args, flat_arg_fake_tensors) = self.validate_and_convert_non_fake_tensors(
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/_subclasses/fake_tensor.py", line 2779, in validate_and_convert_non_fake_tensors
[rank0]:     validated_args = [validate(a) for a in flat_args]
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/_subclasses/fake_tensor.py", line 2779, in <listcomp>
[rank0]:     validated_args = [validate(a) for a in flat_args]
[rank0]:   File "/data/users/hirsheybar/a/pytorch/torch/_subclasses/fake_tensor.py", line 2767, in validate
[rank0]:     raise AssertionError(
[rank0]: torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function <built-in method randint_like of type object at 0x7f5115b45820>(*(DTensor(local_tensor=FakeTensor(..., device='cuda:0', size=(64,)), device_mesh=DeviceMesh('cuda', [0, 1]), placements=(Shard(dim=0),)), 0, 65536), **{'dtype': torch.int32}): got AssertionError("Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.to.dtype_layout(tensor([...], size=(16,), dtype=torch.uint8), dtype=torch.uint8, layout=torch.strided, device=device(type='cuda', index=0))")