Post-training suffering from CUDA error

[youthHan]

Thanks for releasing the training codes and pipeline. While I'm trying to reproduce the libero-long results, I encountered CUDA errors and had to wrap with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION): on

molmoact/olmo/nn/image_vit.py

Line 215 in 8895cba

attn_output = F.scaled_dot_product_attention(

However, this will greatly slow down the training, for 80000 steps, it currently needs 21 days to finish.

I tried different dockers (including the one in this repo), cu124+torch2.6 and cu126+torch2.7. All these trials result in CUDA errors. Could anyone that successfully starts training share there libs and versions?

[youthHan]

BTW, I'm using 8xH100 GPUs.

[youthHan]

[2025-09-14 08:28:15] INFO [olmo.train.trainer:852, rank=0] Pre-train system metrics System/Peak GPU Memory (MB)=65,145 [rank0]:[W914 08:28:31.376698718 Module.cpp:182] symbolizing C++ stack trace for exception; if this hangs, rerun with TORCH_DISABLE_ADDR2LINE=1... [2025-09-14 08:28:32] WARNING [py.warnings:110, rank=0] /opt/conda/lib/python3.11/site-packages/torch/autograd/graph.py:823: UserWarning: Error detected in ScaledDotProductFlashAttentionBackward0. Traceback of forward call that caused the error: File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/launch_scripts/train_multitask_model.py", line 541, in run_trainer(cfg) File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/scripts/train.py", line 466, in run_trainer trainer.fit() File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 1097, in fit metrics = self.train_step(batch, compute_metrics=should_log_this_step) File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 745, in train_step loss = self.train_batch(batch, True) File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 704, in train_batch ce_loss, z_loss, model_out = self.model_forward( File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 671, in model_forward model_out = self.fsdp_model( File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in wrapped_call_impl return self.call_impl(*args, kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl return forward_call(args, kwargs) File "/opt/conda/lib/python3.11/site-packages/peft/peft_model.py", line 881, in forward return self.get_base_model()(args, kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl return self._call_impl(args, kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl return inner() File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1793, in inner result = forward_call(args, kwargs) File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/models/molmo/molmo.py", line 428, in forward image_features = self.vision_backbone(images, image_masks, pooled_patches_idx) File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl return self._call_impl(args, kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl return inner() File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1793, in inner result = forward_call(args, **kwargs) File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/nn/vision_backbone.py", line 362, in forward pooled_features = self.image_pooling_2d(query, to_pool, attn_mask=attn_mask) File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1737, in _wrapped_call_impl return self._compiled_call_impl(args, **kwargs) # type: ignore[misc] File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl return inner() File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1793, in inner result = forward_call(args, *kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py", line 170, in forward return self.checkpoint_fn( # type: ignore[misc] File "/opt/conda/lib/python3.11/site-packages/torch/_compile.py", line 32, in inner return disable_fn(args, kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 745, in fn return fn(*args, **kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/utils/checkpoint.py", line 496, in checkpoint ret = function(*args, *kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl return self.call_impl(*args, kwargs) File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl return forward_call(args, kwargs) File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/nn/image_vit.py", line 215, in forward attn_output = F.scaled_dot_product_attention( (Triggered internally at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:122.) return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass [rank0]:[W914 08:28:32.170774890 Module.cpp:182] symbolizing C++ stack trace for exception; if this hangs, rerun with TORCH_DISABLE_ADDR2LINE=1... Traceback (most recent call last): File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/scripts/train.py", line 466, in run_trainer trainer.fit() File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 1097, in fit metrics = self.train_step(batch, compute_metrics=should_log_this_step) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 745, in train_step loss = self.train_batch(batch, True) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 728, in train_batch loss.backward() File "/opt/conda/lib/python3.11/site-packages/torch/_tensor.py", line 626, in backward torch.autograd.backward( File "/opt/conda/lib/python3.11/site-packages/torch/autograd/init.py", line 347, in backward _engine_run_backward( File "/opt/conda/lib/python3.11/site-packages/torch/autograd/graph.py", line 823, in _engine_run_backward return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RuntimeError: CUDA error: an illegal memory access was encountered Compile with TORCH_USE_CUDA_DSA to enable device-side assertions. Exception raised from c10_cuda_check_implementation at /pytorch/c10/cuda/CUDAException.cpp:43 (most recent call first): C++ CapturedTraceback: #4 std::_Function_handler<std::shared_ptr<c10::LazyValuestd::string const> (), c10::SetStackTraceFetcher(std::function<std::string ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&) from Logging.cpp:0 #5 c10::Error::Error(c10::SourceLocation, std::string) from ??:0 #6 c10::detail::torchCheckFail(char const, char const, unsigned int, std::string const&) from ??:0 #7 c10::cuda::c10_cuda_check_implementation(int, char const, char const, int, bool) from ??:0 #8 void pytorch_flash::run_flash_bwd_seqk_parallel<pytorch_flash::Flash_bwd_kernel_traits<96, 64, 128, 8, 2, 4, 4, true, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<96, 64, 128, 8, cutlass::bfloat16_t> >, false>(pytorch_flash::Flash_bwd_params&, CUstream_st) from ??:0 #9 void pytorch_flash::run_mha_bwd<cutlass::bfloat16_t, 96>(pytorch_flash::Flash_bwd_params&, CUstream_st) from ??:0 #10 pytorch_flash::mha_bwd(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, std::optionalat::Tensor&, std::optionalat::Tensor&, std::optionalat::Tensor&, std::optionalat::Tensor&, float, float, bool, int, int, bool, at::Tensor, at::Tensor) from :0 #11 at::native::flash_attention_backward(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, long, long, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optional, std::optional) from ??:0 #12 at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA___flash_attention_backward(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt) from RegisterCUDA.cpp:0 #13 c10::impl::wrap_kernel_functor_unboxed<c10::impl::detail::WrapFunctionIntoFunctor<c10::CompileTimeFunctionPointer<std::tuple<at::Tensor, at::Tensor, at::Tensor> (at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA___flash_attention_backward>, std::tuple<at::Tensor, at::Tensor, at::Tensor>, c10::guts::typelist::typelist<at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt > >, std::tuple<at::Tensor, at::Tensor, at::Tensor> (at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt)>::call(c10::OperatorKernel, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt) from RegisterCUDA.cpp:0 #14 std::tuple<at::Tensor, at::Tensor, at::Tensor> c10::callUnboxedKernelFunction<std::tuple<at::Tensor, at::Tensor, at::Tensor>, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt >(void, c10::OperatorKernel, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt&&, c10::SymInt&&, double&&, bool&&, at::Tensor const&, at::Tensor const&, std::optional&&, std::optionalc10::SymInt&&, std::optionalc10::SymInt&&) from :0 #15 at::_ops::_flash_attention_backward::call(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt) from ??:0 #16 at::native::scaled_dot_product_flash_attention_backward_cuda(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, long, long, double, bool, at::Tensor const&, at::Tensor const&, std::optional) from ??:0 #17 at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA___scaled_dot_product_flash_attention_backward(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional) from RegisterCUDA.cpp:0 #18 c10::impl::make_boxed_from_unboxed_functor<c10::impl::detail::WrapFunctionIntoFunctor<c10::CompileTimeFunctionPointer<std::tuple<at::Tensor, at::Tensor, at::Tensor> (at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA___scaled_dot_product_flash_attention_backward>, std::tuple<at::Tensor, at::Tensor, at::Tensor>, c10::guts::typelist::typelist<at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional > >, false>::call(c10::OperatorKernel, c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocatorc10::IValue >) from RegisterCUDA.cpp:0 #19 c10::OperatorHandle::redispatchBoxed(c10::DispatchKeySet, std::vector<c10::IValue, std::allocatorc10::IValue >) const from :0 #20 torch::autograd::autogradNotImplementedFallbackImpl(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocatorc10::IValue >) from autograd_not_implemented_fallback.cpp:0 #21 c10::impl::BoxedKernelWrapper<std::tuple<at::Tensor, at::Tensor, at::Tensor> (at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional), void>::call(c10::BoxedKernel const&, c10::OperatorHandle const&, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional) from :0 #22 at::_ops::_scaled_dot_product_flash_attention_backward::call(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional) from ??:0 #23 torch::autograd::generated::ScaledDotProductFlashAttentionBackward0::apply(std::vector<at::Tensor, std::allocatorat::Tensor >&&) from ??:0 #24 torch::autograd::Node::operator()(std::vector<at::Tensor, std::allocatorat::Tensor >&&) from :0 #25 torch::autograd::Engine::evaluate_function(std::shared_ptrtorch::autograd::GraphTask&, torch::autograd::Node, torch::autograd::InputBuffer&, std::shared_ptrtorch::autograd::ReadyQueue const&) from ??:0 #26 torch::autograd::Engine::thread_main(std::shared_ptrtorch::autograd::GraphTask const&) from ??:0 #27 torch::autograd::Engine::thread_init(int, std::shared_ptrtorch::autograd::ReadyQueue const&, bool) from ??:0 #28 torch::autograd::python::PythonEngine::thread_init(int, std::shared_ptrtorch::autograd::ReadyQueue const&, bool) from ??:0 #29 execute_native_thread_routine from thread48.o:0 #30 pthread_condattr_setpshared from ??:0 #31 __xmknodat from ??:0 During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/launch_scripts/train_multitask_model.py", line 541, in run_trainer(cfg) File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/scripts/train.py", line 376, in run_trainer with Trainer( File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 1335, in exit self.close(0 if exc_type is None else 1) File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py", line 1319, in close gc_cuda() File "/mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/torch_util.py", line 167, in gc_cuda torch.cuda.empty_cache() File "/opt/conda/lib/python3.11/site-packages/torch/cuda/memory.py", line 218, in empty_cache torch._C._cuda_emptyCache() RuntimeError: CUDA error: an illegal memory access was encountered Compile with TORCH_USE_CUDA_DSA to enable device-side assertions. Exception raised from c10_cuda_check_implementation at /pytorch/c10/cuda/CUDAException.cpp:43 (most recent call first): C++ CapturedTraceback: #4 std::_Function_handler<std::shared_ptr<c10::LazyValuestd::string const> (), c10::SetStackTraceFetcher(std::function<std::string ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&) from Logging.cpp:0 #5 c10::Error::Error(c10::SourceLocation, std::string) from ??:0 #6 c10::detail::torchCheckFail(char const, char const, unsigned int, std::string const&) from ??:0 #7 c10::cuda::c10_cuda_check_implementation(int, char const, char const, int, bool) from ??:0 #8 c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::release_block(c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::Block, std::shared_ptrc10::GatheredContext const&) from CUDACachingAllocator.cpp:0 #9 c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::release_blocks(c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::BlockPool&, std::shared_ptrc10::GatheredContext const&) from CUDACachingAllocator.cpp:0 #10 c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::emptyCache() from :0 #11 THCPModule_emptyCache(_object, _object) from ??:0 #12 cfunction_vectorcall_NOARGS from /usr/local/src/conda/python-3.11.13/Objects/methodobject.c:486 #13 _PyObject_VectorcallTstate from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_call.h:92 #14 _PyEval_EvalFrameDefault from /usr/local/src/conda/python-3.11.13/Python/ceval.c:4769 #15 _PyEval_EvalFrame from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_ceval.h:73 #16 method_vectorcall from /usr/local/src/conda/python-3.11.13/Objects/classobject.c:59 #17 _PyObject_VectorcallTstate from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_call.h:92 #18 _PyEval_EvalFrameDefault from /usr/local/src/conda/python-3.11.13/Python/ceval.c:4451 #19 _PyEval_EvalFrame from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_ceval.h:73 #20 PyEval_EvalCode from /usr/local/src/conda/python-3.11.13/Python/ceval.c:1148 #21 run_eval_code_obj from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1741 #22 run_mod from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1762 #23 pyrun_file from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1657 #24 _PyRun_SimpleFileObject from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:440 #25 _PyRun_AnyFileObject from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:79 #26 pymain_run_file_obj from /usr/local/src/conda/python-3.11.13/Modules/main.c:360 #27 Py_BytesMain from /usr/local/src/conda/python-3.11.13/Modules/main.c:738 #28 __libc_init_first from ??:0 #29 __libc_start_main from ??:0 #30 _start from ??:0 [2025-09-14 08:28:50] CRITICAL [olmo.util:177, rank=0] Uncaught RuntimeError: CUDA error: an illegal memory access was encountered Compile with TORCH_USE_CUDA_DSA to enable device-side assertions. Exception raised from c10_cuda_check_implementation at /pytorch/c10/cuda/CUDAException.cpp:43 (most recent call first): C++ CapturedTraceback: #4 std::_Function_handler<std::shared_ptr<c10::LazyValuestd::string const> (), c10::SetStackTraceFetcher(std::function<std::string ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&) from Logging.cpp:0 #5 c10::Error::Error(c10::SourceLocation, std::string) from ??:0 #6 c10::detail::torchCheckFail(char const, char const, unsigned int, std::string const&) from ??:0 #7 c10::cuda::c10_cuda_check_implementation(int, char const, char const, int, bool) from ??:0 #8 c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::release_block(c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::Block, std::shared_ptrc10::GatheredContext const&) from CUDACachingAllocator.cpp:0 #9 c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::release_blocks(c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::BlockPool&, std::shared_ptrc10::GatheredContext const&) from CUDACachingAllocator.cpp:0 #10 c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::emptyCache() from :0 #11 THCPModule_emptyCache(_object, _object) from ??:0 #12 cfunction_vectorcall_NOARGS from /usr/local/src/conda/python-3.11.13/Objects/methodobject.c:486 #13 _PyObject_VectorcallTstate from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_call.h:92 #14 _PyEval_EvalFrameDefault from /usr/local/src/conda/python-3.11.13/Python/ceval.c:4769 #15 _PyEval_EvalFrame from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_ceval.h:73 #16 method_vectorcall from /usr/local/src/conda/python-3.11.13/Objects/classobject.c:59 #17 _PyObject_VectorcallTstate from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_call.h:92 #18 _PyEval_EvalFrameDefault from /usr/local/src/conda/python-3.11.13/Python/ceval.c:4451 #19 _PyEval_EvalFrame from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_ceval.h:73 #20 PyEval_EvalCode from /usr/local/src/conda/python-3.11.13/Python/ceval.c:1148 #21 run_eval_code_obj from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1741 #22 run_mod from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1762 #23 pyrun_file from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1657 #24 _PyRun_SimpleFileObject from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:440 #25 _PyRun_AnyFileObject from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:79 #26 pymain_run_file_obj from /usr/local/src/conda/python-3.11.13/Modules/main.c:360 #27 Py_BytesMain from /usr/local/src/conda/python-3.11.13/Modules/main.c:738 #28 __libc_init_first from ??:0 #29 __libc_start_main from ??:0 #30 _start from ??:0 ╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/scripts/train.py:466 in run_trainer │ │ │ │ 463 │ │ # Ready to start training │ │ 464 │ │ if not cfg.dry_run: │ │ 465 │ │ │ log.info("Starting training...") │ │ ❱ 466 │ │ │ trainer.fit() │ │ 467 │ │ │ log.info("Training complete") │ │ 468 │ │ else: │ │ 469 │ │ │ log.info("Dry run complete") │ │ │ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py:1097 in fit │ │ │ │ 1094 │ │ │ │ │ │ self._train_start_time = time.monotonic() │ │ 1095 │ │ │ │ │ │ │ 1096 │ │ │ │ │ # Run train step on batch. │ │ ❱ 1097 │ │ │ │ │ metrics = self.train_step(batch, compute_metrics=should_log_this_step │ │ 1098 │ │ │ │ │ │ │ 1099 │ │ │ │ │ # Maybe collect other metrics. │ │ 1100 │ │ │ │ │ if should_log_this_step: │ │ │ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py:745 in train_step │ │ │ │ 742 │ │ # Run forward-backward pass │ │ 743 │ │ if compute_metrics: │ │ 744 │ │ │ self._train_metrics.reset() │ │ ❱ 745 │ │ loss = self.train_batch(batch, True) │ │ 746 │ │ │ │ 747 │ │ if compute_metrics: │ │ 748 │ │ │ metrics = {f"train/{k}": v for k, v in self._train_metrics.compute().items()} │ │ │ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py:728 in train_batch │ │ │ │ 725 │ │ │ del model_out │ │ 726 │ │ │ │ │ 727 │ │ │ # Run backward pass. │ │ ❱ 728 │ │ │ loss.backward() │ │ 729 │ │ │ total_loss += loss.detach() │ │ 730 │ │ │ │ 731 │ │ return total_loss │ │ │ │ /opt/conda/lib/python3.11/site-packages/torch/_tensor.py:626 in backward │ │ │ │ 623 │ │ │ │ create_graph=create_graph, │ │ 624 │ │ │ │ inputs=inputs, │ │ 625 │ │ │ ) │ │ ❱ 626 │ │ torch.autograd.backward( │ │ 627 │ │ │ self, gradient, retain_graph, create_graph, inputs=inputs │ │ 628 │ │ ) │ │ 629 │ │ │ │ /opt/conda/lib/python3.11/site-packages/torch/autograd/init.py:347 in backward │ │ │ │ 344 │ # The reason we repeat the same comment below is that │ │ 345 │ # some Python versions print out the first line of a multi-line function │ │ 346 │ # calls in the traceback and some print out the last line │ │ ❱ 347 │ engine_run_backward( │ │ 348 │ │ tensors, │ │ 349 │ │ grad_tensors, │ │ 350 │ │ retain_graph, │ │ │ │ /opt/conda/lib/python3.11/site-packages/torch/autograd/graph.py:823 in _engine_run_backward │ │ │ │ 820 │ if attach_logging_hooks: │ │ 821 │ │ unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs) │ │ 822 │ try: │ │ ❱ 823 │ │ return Variable._execution_engine.run_backward( # Calls into the C++ engine to ru │ │ 824 │ │ │ t_outputs, args, kwargs │ │ 825 │ │ ) # Calls into the C++ engine to run the backward pass │ │ 826 │ finally: │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ RuntimeError: CUDA error: an illegal memory access was encountered Compile with TORCH_USE_CUDA_DSA to enable device-side assertions. Exception raised from c10_cuda_check_implementation at /pytorch/c10/cuda/CUDAException.cpp:43 (most recent call first): C++ CapturedTraceback: #4 std::_Function_handler<std::shared_ptr<c10::LazyValuestd::string const> (), c10::SetStackTraceFetcher(std::function<std::string ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&) from Logging.cpp:0 #5 c10::Error::Error(c10::SourceLocation, std::string) from ??:0 #6 c10::detail::torchCheckFail(char const, char const, unsigned int, std::string const&) from ??:0 #7 c10::cuda::c10_cuda_check_implementation(int, char const, char const, int, bool) from ??:0 #8 void pytorch_flash::run_flash_bwd_seqk_parallel<pytorch_flash::Flash_bwd_kernel_traits<96, 64, 128, 8, 2, 4, 4, true, false, cutlass::bfloat16_t, pytorch_flash::Flash_kernel_traits<96, 64, 128, 8, cutlass::bfloat16_t> >, false>(pytorch_flash::Flash_bwd_params&, CUstream_st) from ??:0 #9 void pytorch_flash::run_mha_bwd<cutlass::bfloat16_t, 96>(pytorch_flash::Flash_bwd_params&, CUstream_st*) from ??:0 #10 pytorch_flash::mha_bwd(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, std::optionalat::Tensor&, std::optionalat::Tensor&, std::optionalat::Tensor&, std::optionalat::Tensor&, float, float, bool, int, int, bool, at::Tensor, at::Tensor) from :0 #11 at::native::flash_attention_backward(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, long, long, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optional, std::optional) from ??:0 #12 at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA___flash_attention_backward(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt) from RegisterCUDA.cpp:0 #13 c10::impl::wrap_kernel_functor_unboxed<c10::impl::detail::WrapFunctionIntoFunctor<c10::CompileTimeFunctionPointer<std::tuple<at::Tensor, at::Tensor, at::Tensor> (at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA___flash_attention_backward>, std::tuple<at::Tensor, at::Tensor, at::Tensor>, c10::guts::typelist::typelist<at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt > >, std::tuple<at::Tensor, at::Tensor, at::Tensor> (at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt) from RegisterCUDA.cpp:0 #14 std::tuple<at::Tensor, at::Tensor, at::Tensor> c10::callUnboxedKernelFunction<std::tuple<at::Tensor, at::Tensor, at::Tensor>, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt >(void*, c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt&&, c10::SymInt&&, double&&, bool&&, at::Tensor const&, at::Tensor const&, std::optional&&, std::optionalc10::SymInt&&, std::optionalc10::SymInt&&) from :0 #15 at::_ops::_flash_attention_backward::call(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional, std::optionalc10::SymInt, std::optionalc10::SymInt) from ??:0 #16 at::native::scaled_dot_product_flash_attention_backward_cuda(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, long, long, double, bool, at::Tensor const&, at::Tensor const&, std::optional) from ??:0 #17 at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA___scaled_dot_product_flash_attention_backward(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional) from RegisterCUDA.cpp:0 #18 c10::impl::make_boxed_from_unboxed_functor<c10::impl::detail::WrapFunctionIntoFunctor<c10::CompileTimeFunctionPointer<std::tuple<at::Tensor, at::Tensor, at::Tensor> (at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA___scaled_dot_product_flash_attention_backward>, std::tuple<at::Tensor, at::Tensor, at::Tensor>, c10::guts::typelist::typelist<at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional > >, false>::call(c10::OperatorKernel*, c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocatorc10::IValue >) from RegisterCUDA.cpp:0 #19 c10::OperatorHandle::redispatchBoxed(c10::DispatchKeySet, std::vector<c10::IValue, std::allocatorc10::IValue >) const from :0 #20 torch::autograd::autogradNotImplementedFallbackImpl(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocatorc10::IValue >) from autograd_not_implemented_fallback.cpp:0 #21 c10::impl::BoxedKernelWrapper<std::tuple<at::Tensor, at::Tensor, at::Tensor> (at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional), void>::call(c10::BoxedKernel const&, c10::OperatorHandle const&, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional) from :0 #22 at::_ops::_scaled_dot_product_flash_attention_backward::call(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::SymInt, c10::SymInt, double, bool, at::Tensor const&, at::Tensor const&, std::optional) from ??:0 #23 torch::autograd::generated::ScaledDotProductFlashAttentionBackward0::apply(std::vector<at::Tensor, std::allocatorat::Tensor >&&) from ??:0 #24 torch::autograd::Node::operator()(std::vector<at::Tensor, std::allocatorat::Tensor >&&) from :0 #25 torch::autograd::Engine::evaluate_function(std::shared_ptrtorch::autograd::GraphTask&, torch::autograd::Node, torch::autograd::InputBuffer&, std::shared_ptrtorch::autograd::ReadyQueue const&) from ??:0 #26 torch::autograd::Engine::thread_main(std::shared_ptrtorch::autograd::GraphTask const&) from ??:0 #27 torch::autograd::Engine::thread_init(int, std::shared_ptrtorch::autograd::ReadyQueue const&, bool) from ??:0 #28 torch::autograd::python::PythonEngine::thread_init(int, std::shared_ptrtorch::autograd::ReadyQueue const&, bool) from ??:0 #29 execute_native_thread_routine from thread48.o:0 #30 pthread_condattr_setpshared from ??:0 #31 __xmknodat from ??:0 During handling of the above exception, another exception occurred: ╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/launch_scripts/train_multitask_model.py: │ │ │ │ 538 │ conf = OmegaConf.create(cfg) │ │ 539 │ conf.merge_with_dotlist([clean_opt(arg) for arg in other_args]) │ │ 540 │ cfg = cast(TrainConfig, OmegaConf.to_object(conf)) │ │ ❱ 541 │ run_trainer(cfg) │ │ 542 │ │ │ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/scripts/train.py:376 in run_trainer │ │ │ │ 373 │ │ │ │ │ OmegaConf.to_yaml(cfg, resolve=True), cfg.save_overwrite) │ │ 374 │ barrier() │ │ 375 │ │ │ ❱ 376 │ with Trainer( │ │ 377 │ │ cfg=cfg, │ │ 378 │ │ epoch=cfg.epoch, │ │ 379 │ │ model=olmo_model, │ │ │ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py:1335 in exit │ │ │ │ 1332 │ │ │ 1333 │ def exit(self, exc_type, exc_val, exc_tb) -> None: │ │ 1334 │ │ del exc_val, exc_tb │ │ ❱ 1335 │ │ self.close(0 if exc_type is None else 1) │ │ 1336 │ │ │ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/train/trainer.py:1319 in close │ │ │ │ 1316 │ │ │ 1317 │ │ │ 1318 │ def close(self, exit_code: int = 0) -> None: │ │ ❱ 1319 │ │ gc_cuda() │ │ 1320 │ │ │ │ 1321 │ │ if self._gc_init_state: │ │ 1322 │ │ │ gc.enable() │ │ │ │ /mnt/bn/kinetics-lp-maliva/playground_projects/MolmoAct/olmo/torch_util.py:167 in gc_cuda │ │ │ │ 164 def gc_cuda(): │ │ 165 │ gc.collect() │ │ 166 │ if torch.cuda.is_available(): │ │ ❱ 167 │ │ torch.cuda.empty_cache() │ │ 168 │ │ 169 │ │ 170 def listinstr(lst, s, delimiter=None): │ │ │ │ /opt/conda/lib/python3.11/site-packages/torch/cuda/memory.py:218 in empty_cache │ │ │ │ 215 │ │ more details about GPU memory management. │ │ 216 │ """ │ │ 217 │ if is_initialized(): │ │ ❱ 218 │ │ torch._C._cuda_emptyCache() │ │ 219 │ │ 220 │ │ 221 def memory_stats(device: Union[Device, int] = None) -> Dict[str, Any]: │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ RuntimeError: CUDA error: an illegal memory access was encountered Compile with TORCH_USE_CUDA_DSA to enable device-side assertions. Exception raised from c10_cuda_check_implementation at /pytorch/c10/cuda/CUDAException.cpp:43 (most recent call first): C++ CapturedTraceback: #4 std::_Function_handler<std::shared_ptr<c10::LazyValuestd::string const> (), c10::SetStackTraceFetcher(std::function<std::string ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&) from Logging.cpp:0 #5 c10::Error::Error(c10::SourceLocation, std::string) from ??:0 #6 c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) from ??:0 #7 c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) from ??:0 #8 c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::release_block(c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::Block*, std::shared_ptrc10::GatheredContext const&) from CUDACachingAllocator.cpp:0 #9 c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::release_blocks(c10::cuda::CUDACachingAllocator::Native::(anonymous namespace)::BlockPool&, std::shared_ptrc10::GatheredContext const&) from CUDACachingAllocator.cpp:0 #10 c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::emptyCache() from :0 #11 THCPModule_emptyCache(_object*, _object*) from ??:0 #12 cfunction_vectorcall_NOARGS from /usr/local/src/conda/python-3.11.13/Objects/methodobject.c:486 #13 _PyObject_VectorcallTstate from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_call.h:92 #14 _PyEval_EvalFrameDefault from /usr/local/src/conda/python-3.11.13/Python/ceval.c:4769 #15 _PyEval_EvalFrame from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_ceval.h:73 #16 method_vectorcall from /usr/local/src/conda/python-3.11.13/Objects/classobject.c:59 #17 _PyObject_VectorcallTstate from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_call.h:92 #18 _PyEval_EvalFrameDefault from /usr/local/src/conda/python-3.11.13/Python/ceval.c:4451 #19 _PyEval_EvalFrame from /usr/local/src/conda/python-3.11.13/Include/internal/pycore_ceval.h:73 #20 PyEval_EvalCode from /usr/local/src/conda/python-3.11.13/Python/ceval.c:1148 #21 run_eval_code_obj from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1741 #22 run_mod from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1762 #23 pyrun_file from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:1657 #24 _PyRun_SimpleFileObject from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:440 #25 _PyRun_AnyFileObject from /usr/local/src/conda/python-3.11.13/Python/pythonrun.c:79 #26 pymain_run_file_obj from /usr/local/src/conda/python-3.11.13/Modules/main.c:360 #27 Py_BytesMain from /usr/local/src/conda/python-3.11.13/Modules/main.c:738 #28 __libc_init_first from ??:0 #29 __libc_start_main from ??:0 #30 _start from ??:0 wandb: wandb: 🚀 View run replicate-libero-long at: https://wandb.ai/web-videos-for-robotics/flexvla/runs/8zcmmm5q E0914 08:28:56.277000 431806 site-packages/torch/distributed/elastic/multiprocessing/api.py:869] failed (exitcode: 1) local_rank: 0 (pid: 431854) of binary: /opt/conda/bin/python Traceback (most recent call last): File "/opt/conda/bin/torchrun", line 7, in sys.exit(main()) ^^^^^^ File "/opt/conda/lib/python3.11/site-packages/torch/distributed/elastic/multiprocessing/errors/init.py", line 355, in wrapper return f(*args, **kwargs) ^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.11/site-packages/torch/distributed/run.py", line 918, in main run(args) File "/opt/conda/lib/python3.11/site-packages/torch/distributed/run.py", line 909, in run elastic_launch( File "/opt/conda/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 138, in call return launch_agent(self._config, self._entrypoint, list(args)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 269, in launch_agent raise ChildFailedError( torch.distributed.elastic.multiprocessing.errors.ChildFailedError: ============================================================ launch_scripts/train_multitask_model.py FAILED ------------------------------------------------------------ Failures: <NO_OTHER_FAILURES> ------------------------------------------------------------ Root Cause (first observed failure): [0]: time : 2025-09-14_08:28:56 host : g340-cda2-1230-a624-c354-6233-629b.byted.org rank : 0 (local_rank: 0) exitcode : 1 (pid: 431854) error_file: <N/A> traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html