Add demos/dimensionality_ladder

.gitattributes

@@ -202,3 +202,10 @@ demos/dataset_trajectories/traj_expert_good.png filter=lfs diff=lfs merge=lfs -t
 demos/dataset_trajectories/traj_expert_suboptimal.png filter=lfs diff=lfs merge=lfs -text
 demos/dataset_trajectories/traj_policy_failed.png filter=lfs diff=lfs merge=lfs -text
 demos/dataset_trajectories/traj_policy_good.png filter=lfs diff=lfs merge=lfs -text
+demos/dimensionality_ladder/dim_curve.png filter=lfs diff=lfs merge=lfs -text
+demos/dimensionality_ladder/probe10k_d1024.png filter=lfs diff=lfs merge=lfs -text
+demos/dimensionality_ladder/probe10k_d512.png filter=lfs diff=lfs merge=lfs -text
+demos/dimensionality_ladder/probe10k_d64.png filter=lfs diff=lfs merge=lfs -text
+demos/dimensionality_ladder/probe_d1024.png filter=lfs diff=lfs merge=lfs -text
+demos/dimensionality_ladder/probe_d512.png filter=lfs diff=lfs merge=lfs -text
+demos/dimensionality_ladder/probe_d64.png filter=lfs diff=lfs merge=lfs -text

demos/dimensionality_ladder/probe10k_d1024/decoder_probe.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d353b56395a8afb84328cda555cb7e1acaf8705e851d40b41a09f95dc449ba35
+size 5738957

demos/dimensionality_ladder/probe10k_d512/decoder_probe.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eed8078569177b40ce0cc06ab91607db949fbb6d0b0a75621fad58b3fc8386e9
+size 5743263

demos/dimensionality_ladder/probe10k_d64/decoder_probe.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f665d870fbc3c946b252e2f3a7fa404ada98b5a5bdafd0cda036f36f8fe75b9b
+size 5724134

demos/dimensionality_ladder/probe_d1024/decoder_probe.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18840445f5cbec557560636edf1fe52bd2ce37aa7cde88233a30441591117aab
+size 5740892

demos/dimensionality_ladder/probe_d512/decoder_probe.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43aa0298dfb7e9af13b8bd894c883a70ea5235ebd71110c5ecb170753e1db930
+size 5743886

demos/dimensionality_ladder/probe_d64/decoder_probe.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d62ccb08b63c9603f9bb9fa1de49b22c8c9da9d4dca23e1b76951122d135b9dc
+size 5728740

demos/dimensionality_ladder/smoke_d1024_r2.log

@@ -0,0 +1,187 @@
+/root/simdist/simdist_uwlab/scripts/train_model.py:14: UserWarning: 
+The version_base parameter is not specified.
+Please specify a compatability version level, or None.
+Will assume defaults for version 1.1
+  @hydra.main(**get_train_model_hydra_config())
+/root/simdist/venv_train/lib/python3.10/site-packages/hydra/_internal/hydra.py:119: UserWarning: Future Hydra versions will no longer change working directory at job runtime by default.
+See https://hydra.cc/docs/1.2/upgrades/1.1_to_1.2/changes_to_job_working_dir/ for more information.
+  ret = run_job(
+INFO:2026-08-13 10:00:36,200:jax._src.xla_bridge:867: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
+Training will stop after 250 steps.
+Running training with config:  {'system': {'name': 'omnireset_ur10e', 'proprio_obs': {'types': [{'name': 'last_gripper_action', 'dim': 1, 'noise': 0.01}, {'name': 'last_arm_action', 'dim': 6, 'noise': 0.01}, {'name': 'arm_joint_pos', 'dim': 6, 'noise': 0.005}, {'name': 'end_effector_pose', 'dim': 6, 'noise': 0.003}, {'name': 'binary_contact', 'dim': 1, 'noise': 0.05}]}, 'extero_obs': {'types': [{'name': 'front_rgb', 'dim': [120, 160, 3], 'noise': 0.01}, {'name': 'side_rgb', 'dim': [120, 160, 3], 'noise': 0.01}, {'name': 'wrist_rgb', 'dim': [120, 160, 3], 'noise': 0.01}]}, 'cmd': {'dim': 0}, 'actions': [{'name': 'ee_delta_x', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_y', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_z', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_rx', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_ry', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_rz', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'gripper_binary', 'min_noise': 0.0, 'max_noise': 0.0, 'flip_prob_min': 0.01, 'flip_prob_max': 0.1}]}, 'model': {'type': 'manipulator_world_model', 'dataset': {'type': 'manipulator_world_model', 'history_length': 5, 'prediction_length': 5, 'jpeg_backend': 'auto', 'augmentations': {'add_noise': True, 'image': {'enabled': True, 'random_crop': {'enabled': True, 'padding': [4, 6]}, 'gaussian_blur': {'enabled': True, 'prob': 0.5, 'kernel_size': 5, 'sigma': [0.1, 1.5]}, 'color_jitter': {'enabled': True, 'brightness': 0.2, 'contrast': 0.2, 'saturation': 0.2, 'hue': 0.02}}}}, 'dropout': {'mlp': 0.2, 'attention': 0.05}, 'latent_dim': 1024, 'encoder': {'mlp_hidden_size_factor': 2, 'proprio_obs_layers': 2, 'extero_obs': {'image_resnet': {'share_trunk': True, 'pretrained': True, 'freeze_batch_norm': True, 'normalize_input': True}}, 'action_layers': 0, 'latent_layers': 2, 'latent_norm': True}, 'embedding': {'mlp_hidden_size_factor': 2, 'future_acts_layers': 0, 'future_cmds_layers': 0, 'reward_layers': 0, 'value_layers': 0}, 'dynamics': {'attention': {'layers': 3, 'heads': 4, 'mlp_hidden_size_factor': 4, 'mask': 'causal'}}, 'reward': {'attention': {'layers': 1, 'heads': 1, 'mlp_hidden_size_factor': 4, 'mask': None}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}, 'value': {'attention': {'layers': 1, 'heads': 1, 'mlp_hidden_size_factor': 4, 'mask': None}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}, 'policy': {'attention': {'layers': 4, 'heads': 8, 'mlp_hidden_size_factor': 4, 'mask': 'causal'}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}}, 'loss': {'type': 'world_model', 'weights': {'latent_dynamics': 1.0, 'reward': 1.0, 'value': 1.0, 'action': 4.0}}, 'heads': {}, 'run_name': 'smoke_d1024_r2', 'data': {'dataset_name': 'simdist_nodr', 'num_train_workers': 8, 'num_test_workers': 4, 'prefetch_factor': 2, 'paper_scale_steps': 25000000}, 'training': {'num_epochs': 100, 'max_steps': 250, 'eval_interval': 200, 'learning_rate': 0.0002, 'warmup_steps': 10000, 'decay_steps': 400000, 'end_learning_rate': 0.0001, 'batch_size': 512, 'training_data_ratio': 0.98, 'seed': 42, 'grad_clip_norm': 1.0}, 'checkpoint': {'enabled': True, 'max_to_keep': None, 'resume_checkpoint': None}, 'eval': {'max_batches': 8, 'scatter_points': 5000, 'episodes': {'enabled': True, 'holdout': 128, 'num_auc_episodes': 512, 'num_curve_episodes': 16, 'match_frac': 0.5, 'curve_stride': 5, 'batch_size': 32, 'success': {'metric': 'episode_return', 'threshold': None}}, 'diagnostics': {'enabled': True, 'dirname': 'diagnostics'}, 'decoder_probe': {'enabled': True, 'hidden_size': 128, 'learning_rate': 0.001, 'steps_per_eval': 200, 'image': {'enabled': True, 'base_channels': 128, 'learning_rate': 0.001, 'batch': 32}}}, 'wandb': {'log': False, 'project': 'simdist_train_model', 'entity': None, 'mode': 'online', 'dir': None}}
+Reserved 17360 chunks (128 whole episodes) for the episode probe; they are in neither the train nor the test split.
+Creating a new model...
+[2026-08-13 10:00:36,200][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
+Run fingerprint: {'dataset/processed_dir': '/root/simdist/simdist_uwlab/datasets/sim/simdist_nodr/processed_data_omnireset_ur10e_H-5_T-5', 'dataset/num_steps': 299985, 'dataset/num_trajectories': 261585, 'dataset/generation_git_sha': "unknown (Command '['git', '-C', '/root/simdist/UWLab', 'rev-parse', 'HEAD']' returned non-zero exit status 128.)", 'dataset/expert_checkpoint': 'model_3700.pt', 'dataset/paper_scale_steps': 25000000, 'dataset/paper_scale_fraction': 0.0119994, 'dataset/raw_camera_names': ['front_rgb', 'side_rgb', 'wrist_rgb'], 'dataset/raw_camera_resolution': [120, 160], 'dataset/raw_clean_interval': [5, 10], 'dataset/raw_created_utc': '2026-08-12T20:28:07.508201+00:00', 'dataset/raw_expert_checkpoint': 'model_3700.pt', 'dataset/raw_expert_prob': 0.5, 'dataset/raw_flip_prob_bounds': [0.01, 0.1], 'dataset/raw_git_sha': "unknown (Command '['git', '-C', '/root/simdist/UWLab', 'rev-parse', 'HEAD']' returned non-zero exit status 128.)", 'dataset/raw_image_encoding': 'JPEG (RGB); decode with PIL, or cv2.imdecode(...)[:, :, ::-1]', 'dataset/raw_insertive_object': 'JigV2', 'dataset/raw_jpeg_quality': 90, 'dataset/raw_noise_interval': [1, 5], 'dataset/raw_num_envs': 64, 'dataset/raw_num_rows': 32768, 'dataset/raw_proprio_dims': [1, 6, 6, 6, 1], 'dataset/raw_proprio_terms': ['last_gripper_action', 'last_arm_action', 'arm_joint_pos', 'end_effector_pose', 'binary_contact'], 'dataset/raw_receptive_object': 'BottomEnclosure', 'dataset/raw_row_order': 'step-major: num_envs consecutive rows per env step', 'dataset/raw_seed': 4242, 'dataset/raw_shard_index': 0, 'dataset/raw_sigma_bounds': [0.05, 1.0], 'dataset/raw_suboptimal_checkpoints': ['model_0.pt', 'model_100.pt', 'model_200.pt', 'model_300.pt', 'model_400.pt', 'model_500.pt', 'model_600.pt', 'model_700.pt', 'model_800.pt', 'model_900.pt', 'model_1000.pt', 'model_1100.pt', 'model_1200.pt', 'model_1300.pt', 'model_1400.pt', 'model_1500.pt', 'model_1600.pt', 'model_1700.pt', 'model_1800.pt', 'model_1900.pt', 'model_2000.pt', 'model_2100.pt', 'model_2200.pt', 'model_2300.pt', 'model_2400.pt', 'model_2500.pt', 'model_2600.pt', 'model_2700.pt', 'model_2800.pt', 'model_2900.pt', 'model_3000.pt', 'model_3100.pt', 'model_3200.pt', 'model_3300.pt', 'model_3400.pt', 'model_3500.pt', 'model_3600.pt'], 'dataset/raw_system': 'omnireset_ur10e', 'dataset/raw_task': 'OmniReset-UR10eLinearGripper-RelCartesianOSC-RGB-DataCollection-v0', 'dataset/raw_num_shards': 10, 'code/simdist_git_sha': None, 'code/simdist_git_dirty': None, 'instance/hostname': 'ebc50e8fee3f', 'instance/platform': 'Linux-6.8.0-90-generic-x86_64-with-glibc2.39', 'instance/python': '3.10.20', 'instance/cpu_count': 32, 'instance/cuda_visible_devices': '', 'instance/gpus': ['NVIDIA GeForce RTX 5090']}
+Parameters to be optimized: 171930185
+Total trainable parameters: 171930185
+Total parameters: 171939845
+Gradient clipping ON at global norm 1.0.
+Starting training
+Starting epoch 1/100
+
+  self.pid = os.fork()
+
+
+  self.pid = os.fork()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/diagnostics/step_200
+[2026-08-13 10:03:28,056][absl][INFO] - [thread=MainThread] Failed to get flag value for EXPERIMENTAL_ORBAX_USE_DISTRIBUTED_PROCESS_ID.
+[2026-08-13 10:03:28,056][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
+[2026-08-13 10:03:28,057][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7b9bdc3c80a0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7b9bdc3c80a0>}).
+[2026-08-13 10:03:28,057][absl][INFO] - orbax-checkpoint version: 0.11.25
+[2026-08-13 10:03:28,057][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2.
+[2026-08-13 10:03:28,058][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 10:03:28,058][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/diagnostics) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 10:03:28,060][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/model_config.yaml) to be a directory.
+[2026-08-13 10:03:28,061][absl][INFO] - Found 0 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2
+[2026-08-13 10:03:28,061][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=None, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7b9bdc4816f0>
+[2026-08-13 10:03:28,061][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 10:03:28,061][absl][INFO] - [process=0] Saving checkpoint at step 200
+[2026-08-13 10:03:28,061][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.
+[2026-08-13 10:03:28,062][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp
+[2026-08-13 10:03:28,062][absl][INFO] - Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786615408062491493, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786615408062491493, "commit_timestamp_nsecs": null, "custom_metadata": {}} to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 10:03:28,063][absl][INFO] - Using ThreadSafeKeyValueSignalingClient
+[2026-08-13 10:03:28,063][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 10:03:28,063][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 10:03:28,063][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 10:03:28,064][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 10:03:28,064][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7b9f8a27e9e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 10:03:28,064][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7b9bdc3c9210>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 10:03:28,064][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 10:03:28,064][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7b9f8a27e9e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 10:03:28,064][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7b9bdc3c8eb0>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 10:03:28,087][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 10:03:28,237][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.150858s
+[2026-08-13 10:03:28,253][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 10:03:28,299][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 10:03:28,476][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.223347s
+[2026-08-13 10:03:28,476][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1 Bytes/s (total gbytes: 655.9 MiB) (time elapsed: 411 milliseconds) (per-host)
+[2026-08-13 10:03:28,476][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.411537s (batch_requests_ready=0.015530s, total_serialization_initiated=0.395736s, others=0.000272s)
+[2026-08-13 10:03:28,489][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 3 Bytes/s (total gbytes: 1.3 GiB) (time elapsed: 418 milliseconds) (per-host)
+[2026-08-13 10:03:28,490][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.418576s (batch_requests_ready=0.166588s, total_serialization_initiated=0.238830s, others=0.013158s)
+[2026-08-13 10:03:28,490][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.427128s (all_items=0.000653s, per_item={'default': '0.00037861', 'opt_state': '0.00027418'}, temp_paths=0.426475)
+[2026-08-13 10:03:28,490][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 10:03:29,035][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 10:03:29,197][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.720455s (commit=0.705476s, array_metadata_write=0.014978s)
+[2026-08-13 10:03:29,199][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 655.9 MiB) (time elapsed: a second) (per-host)
+[2026-08-13 10:03:29,201][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 10:03:29,464][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.973969s (commit=0.956193s, array_metadata_write=0.017777s)
+[2026-08-13 10:03:29,466][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 1.3 GiB) (time elapsed: a second) (per-host)
+[2026-08-13 10:03:29,467][absl][INFO] - Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786615408062491493, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 10:03:29,467][absl][INFO] - Updated Metadata={'item_handlers': {'default': 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler', 'opt_state': 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786615408062491493, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 10:03:29,470][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 10:03:29,473][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.005744s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 10:03:29,474][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/default
+[2026-08-13 10:03:29,476][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 10:03:29,480][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.005353s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 10:03:29,480][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp/opt_state
+[2026-08-13 10:03:29,481][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200
+[2026-08-13 10:03:29,482][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200`.
+[2026-08-13 10:03:29,482][absl][INFO] - Finished synchronous save in 1.42 seconds to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200
+[2026-08-13 10:03:29,483][absl][INFO] - [process=0][thread=MainThread][step=200] CheckpointManager Save Finalize is syncing with other hosts...
+[2026-08-13 10:03:29,483][absl][INFO] - [process=0][thread=MainThread][step=200] CheckpointManager Save Finalize is done on all hosts.
+[2026-08-13 10:03:29,483][absl][INFO] - [process=0][thread=MainThread][step=200] Finished synchronous save.
+[2026-08-13 10:03:29,483][absl][INFO] - {'step': 200, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786615408.0618272, 'wait_for_prev_duration_secs': 3.838539123535156e-05, 'checkpointer_blocking_start_time': 1786615408.0619526, 'checkpointer_blocking_duration_secs': 1.421182632446289, 'get_old_steps_start_time': 1786615409.4831524, 'get_old_steps_duration_secs': 9.298324584960938e-05, 'checkpoint_manager_blocking_start_time': 1786615408.0617726, 'checkpoint_manager_blocking_duration_secs': 1.4217627048492432}
+[2026-08-13 10:03:29,483][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 10:03:29,483][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=135928726210368 count=1 at 0x7b9e60275ac0>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7b9bdc481d20>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7b9bdc481d50>, _write_futures=[])
+
+
+
+  self.pid = os.fork()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/diagnostics/step_250
+[2026-08-13 10:04:06,245][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
+[2026-08-13 10:04:06,245][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7b9bbc6bf670>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7b9bbc6bf670>}).
+[2026-08-13 10:04:06,246][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2.
+[2026-08-13 10:04:06,247][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 10:04:06,247][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/diagnostics) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 10:04:06,247][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/200) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 10:04:06,249][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/model_config.yaml) to be a directory.
+[2026-08-13 10:04:06,250][absl][INFO] - Missing metrics for step 200
+[2026-08-13 10:04:06,251][absl][INFO] - Found 1 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2
+[2026-08-13 10:04:06,251][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=None, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7b9bbc6bc070>
+[2026-08-13 10:04:06,251][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 10:04:06,251][absl][INFO] - [process=0] Saving checkpoint at step 250
+[2026-08-13 10:04:06,251][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.
+[2026-08-13 10:04:06,251][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp
+[2026-08-13 10:04:06,252][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 10:04:06,253][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 10:04:06,253][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 10:04:06,253][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 10:04:06,254][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7b9f8a27e9e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 10:04:06,254][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7b9bbc6bed10>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 10:04:06,254][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 10:04:06,254][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7b9f8a27e9e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 10:04:06,254][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7b9bbc6bc910>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 10:04:06,276][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 10:04:06,295][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.018564s
+[2026-08-13 10:04:06,316][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 10 Bytes/s (total gbytes: 655.9 MiB) (time elapsed: 61 milliseconds) (per-host)
+[2026-08-13 10:04:06,319][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.064663s (batch_requests_ready=0.015869s, total_serialization_initiated=0.043807s, others=0.004987s)
+[2026-08-13 10:04:06,320][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 10:04:06,368][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 10:04:06,450][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.130827s
+[2026-08-13 10:04:06,461][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 6 Bytes/s (total gbytes: 1.3 GiB) (time elapsed: 200 milliseconds) (per-host)
+[2026-08-13 10:04:06,464][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.203546s (batch_requests_ready=0.038923s, total_serialization_initiated=0.155796s, others=0.008827s)
+[2026-08-13 10:04:06,465][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.212645s (all_items=0.000286s, per_item={'default': '0.00019908', 'opt_state': '0.00008678'}, temp_paths=0.212359)
+[2026-08-13 10:04:06,465][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 10:04:06,617][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 10:04:06,650][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.333300s (commit=0.318557s, array_metadata_write=0.014743s)
+[2026-08-13 10:04:06,652][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 1 Bytes/s (total gbytes: 655.9 MiB) (time elapsed: 396 milliseconds) (per-host)
+[2026-08-13 10:04:06,652][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 10:04:06,961][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.497463s (commit=0.478020s, array_metadata_write=0.019442s)
+[2026-08-13 10:04:06,962][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 1 Bytes/s (total gbytes: 1.3 GiB) (time elapsed: 701 milliseconds) (per-host)
+[2026-08-13 10:04:06,965][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 10:04:06,969][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.004532s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 10:04:06,969][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/default
+[2026-08-13 10:04:06,972][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 10:04:06,976][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.006225s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 10:04:06,976][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp/opt_state
+[2026-08-13 10:04:06,977][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250
+[2026-08-13 10:04:06,978][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250`.
+[2026-08-13 10:04:06,978][absl][INFO] - Finished synchronous save in 0.73 seconds to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2/250
+[2026-08-13 10:04:06,979][absl][INFO] - [process=0][thread=MainThread][step=250] CheckpointManager Save Finalize is syncing with other hosts...
+[2026-08-13 10:04:06,979][absl][INFO] - [process=0][thread=MainThread][step=250] CheckpointManager Save Finalize is done on all hosts.
+[2026-08-13 10:04:06,979][absl][INFO] - [process=0][thread=MainThread][step=250] Finished synchronous save.
+[2026-08-13 10:04:06,979][absl][INFO] - {'step': 250, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/smoke_d1024_r2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786615446.251424, 'wait_for_prev_duration_secs': 3.62396240234375e-05, 'checkpointer_blocking_start_time': 1786615446.2515323, 'checkpointer_blocking_duration_secs': 0.7273986339569092, 'get_old_steps_start_time': 1786615446.9789448, 'get_old_steps_duration_secs': 7.009506225585938e-05, 'checkpoint_manager_blocking_start_time': 1786615446.25138, 'checkpoint_manager_blocking_duration_secs': 0.7277765274047852}
+[2026-08-13 10:04:06,979][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 10:04:06,979][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=135928726210368 count=1 at 0x7ba05016d040>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7b9bbc6bf9d0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7b9bbc6bf970>, _write_futures=[])
+
+Steps: 250, Metrics: {'steps_per_second': 2.1741483777228665, 'throughput/samples_per_s': 1113.1639693941077, 'lr': 4.980000085197389e-06, 'grad_norm': 28.892051572799684, 'dataloader_wait_frac': 0.003742218742167944, 'gpu_util': 100.0, 'gpu_mem_used_frac': 0.9341552427392891, 'train/loss': 2.801475763320923, 'train/latent_dynamics': 0.5829599499702454, 'train/reward': 0.3898126780986786, 'train/value': 0.6489105224609375, 'train/action': 1.1797924041748047, 'test/loss': 1.7551392316818237, 'test/latent_dynamics': 0.43963271379470825, 'test/reward': 0.2586834728717804, 'test/value': 0.13776764273643494, 'test/action': 0.9190554022789001, 'test/dataloader_wait_frac': 0.13042790410960245, 'train/unweighted/latent_dynamics': 0.5829599499702454, 'train/unweighted/reward': 0.3898126780986786, 'train/unweighted/value': 0.6489105224609375, 'train/unweighted/action': 0.29494810104370117, 'test/unweighted/latent_dynamics': 0.43963271379470825, 'test/unweighted/reward': 0.2586834728717804, 'test/unweighted/value': 0.13776764273643494, 'test/unweighted/action': 0.22976385056972504, 'eval/latent_rollout_error/h1': 0.4513755887746811, 'eval/latent_rollout_error/h2': 0.43578578904271126, 'eval/latent_rollout_error/h3': 0.42806361615657806, 'eval/latent_rollout_error/h4': 0.43942172080278397, 'eval/latent_rollout_error/h5': 0.44238922372460365, 'eval/latent_rollout_error_mean': 0.4394071877002716, 'eval/latent_variance/encoder': 0.0876467414200306, 'eval/latent_variance/predicted': 0.10159797314554453, 'eval/reward_r2': 0.6144998517994975, 'eval/reward_explained_variance': 0.645908110688455, 'eval/reward_pearson_r': 0.8040368767531234, 'eval/value_r2': 0.6901542587893958, 'eval/value_explained_variance': 0.693265143558383, 'eval/value_pearson_r': 0.8329767600976777, 'eval/bc_action_mse': 0.9008812561815147, 'eval/num_batches': 8.0, 'eval/decoder_probe/proprio_mse': 0.5664975047111511, 'eval/decoder_probe/image_mse': 0.020258445292711258, 'eval/value_auc_success': 0.703857421875, 'eval/value_auc_num_episodes': 128.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 1.0, 'epoch': 0}
+Stopping training after 250 steps.

demos/dimensionality_ladder/smoke_d512.log

@@ -0,0 +1,157 @@
+/root/simdist/simdist_uwlab/scripts/train_model.py:14: UserWarning: 
+The version_base parameter is not specified.
+Please specify a compatability version level, or None.
+Will assume defaults for version 1.1
+  @hydra.main(**get_train_model_hydra_config())
+/root/simdist/venv_train/lib/python3.10/site-packages/hydra/_internal/hydra.py:119: UserWarning: Future Hydra versions will no longer change working directory at job runtime by default.
+See https://hydra.cc/docs/1.2/upgrades/1.1_to_1.2/changes_to_job_working_dir/ for more information.
+  ret = run_job(
+INFO:2026-08-13 04:18:53,485:jax._src.xla_bridge:867: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
+Training will stop after 250 steps.
+Running training with config:  {'system': {'name': 'omnireset_ur10e', 'proprio_obs': {'types': [{'name': 'last_gripper_action', 'dim': 1, 'noise': 0.01}, {'name': 'last_arm_action', 'dim': 6, 'noise': 0.01}, {'name': 'arm_joint_pos', 'dim': 6, 'noise': 0.005}, {'name': 'end_effector_pose', 'dim': 6, 'noise': 0.003}, {'name': 'binary_contact', 'dim': 1, 'noise': 0.05}]}, 'extero_obs': {'types': [{'name': 'front_rgb', 'dim': [120, 160, 3], 'noise': 0.01}, {'name': 'side_rgb', 'dim': [120, 160, 3], 'noise': 0.01}, {'name': 'wrist_rgb', 'dim': [120, 160, 3], 'noise': 0.01}]}, 'cmd': {'dim': 0}, 'actions': [{'name': 'ee_delta_x', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_y', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_z', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_rx', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_ry', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_rz', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'gripper_binary', 'min_noise': 0.0, 'max_noise': 0.0, 'flip_prob_min': 0.01, 'flip_prob_max': 0.1}]}, 'model': {'type': 'manipulator_world_model', 'dataset': {'type': 'manipulator_world_model', 'history_length': 5, 'prediction_length': 5, 'jpeg_backend': 'auto', 'augmentations': {'add_noise': True, 'image': {'enabled': True, 'random_crop': {'enabled': True, 'padding': [4, 6]}, 'gaussian_blur': {'enabled': True, 'prob': 0.5, 'kernel_size': 5, 'sigma': [0.1, 1.5]}, 'color_jitter': {'enabled': True, 'brightness': 0.2, 'contrast': 0.2, 'saturation': 0.2, 'hue': 0.02}}}}, 'dropout': {'mlp': 0.2, 'attention': 0.05}, 'latent_dim': 512, 'encoder': {'mlp_hidden_size_factor': 2, 'proprio_obs_layers': 2, 'extero_obs': {'image_resnet': {'share_trunk': True, 'pretrained': True, 'freeze_batch_norm': True, 'normalize_input': True}}, 'action_layers': 0, 'latent_layers': 2, 'latent_norm': True}, 'embedding': {'mlp_hidden_size_factor': 2, 'future_acts_layers': 0, 'future_cmds_layers': 0, 'reward_layers': 0, 'value_layers': 0}, 'dynamics': {'attention': {'layers': 3, 'heads': 4, 'mlp_hidden_size_factor': 4, 'mask': 'causal'}}, 'reward': {'attention': {'layers': 1, 'heads': 1, 'mlp_hidden_size_factor': 4, 'mask': None}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}, 'value': {'attention': {'layers': 1, 'heads': 1, 'mlp_hidden_size_factor': 4, 'mask': None}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}, 'policy': {'attention': {'layers': 4, 'heads': 8, 'mlp_hidden_size_factor': 4, 'mask': 'causal'}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}}, 'loss': {'type': 'world_model', 'weights': {'latent_dynamics': 1.0, 'reward': 1.0, 'value': 1.0, 'action': 4.0}}, 'heads': {}, 'run_name': 'smoke_d512', 'data': {'dataset_name': 'simdist_nodr', 'num_train_workers': 8, 'num_test_workers': 4, 'prefetch_factor': 2, 'paper_scale_steps': 25000000}, 'training': {'num_epochs': 2, 'max_steps': 250, 'eval_interval': 200, 'learning_rate': 0.0002, 'warmup_steps': 10000, 'decay_steps': 400000, 'end_learning_rate': 0.0001, 'batch_size': 512, 'training_data_ratio': 0.998, 'seed': 42, 'grad_clip_norm': 1.0}, 'checkpoint': {'enabled': True, 'max_to_keep': None, 'resume_checkpoint': None}, 'eval': {'max_batches': 8, 'scatter_points': 5000, 'episodes': {'enabled': True, 'holdout': 512, 'num_auc_episodes': 512, 'num_curve_episodes': 16, 'match_frac': 0.5, 'curve_stride': 5, 'batch_size': 32, 'success': {'metric': 'episode_return', 'threshold': None}}, 'diagnostics': {'enabled': True, 'dirname': 'diagnostics'}, 'decoder_probe': {'enabled': True, 'hidden_size': 128, 'learning_rate': 0.001, 'steps_per_eval': 200, 'image': {'enabled': True, 'base_channels': 128, 'learning_rate': 0.001, 'batch': 32}}}, 'wandb': {'log': False, 'project': 'simdist_train_model', 'entity': None, 'mode': 'online', 'dir': None}}
+Reserved 70093 chunks (512 whole episodes) for the episode probe; they are in neither the train nor the test split.
+Creating a new model...
+[2026-08-13 04:18:53,485][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
+Run fingerprint: {'dataset/processed_dir': '/root/simdist/simdist_uwlab/datasets/sim/simdist_nodr/processed_data_omnireset_ur10e_H-5_T-5', 'dataset/num_steps': 299985, 'dataset/num_trajectories': 261585, 'dataset/generation_git_sha': "unknown (Command '['git', '-C', '/root/simdist/UWLab', 'rev-parse', 'HEAD']' returned non-zero exit status 128.)", 'dataset/expert_checkpoint': 'model_3700.pt', 'dataset/paper_scale_steps': 25000000, 'dataset/paper_scale_fraction': 0.0119994, 'dataset/raw_camera_names': ['front_rgb', 'side_rgb', 'wrist_rgb'], 'dataset/raw_camera_resolution': [120, 160], 'dataset/raw_clean_interval': [5, 10], 'dataset/raw_created_utc': '2026-08-12T20:28:07.508201+00:00', 'dataset/raw_expert_checkpoint': 'model_3700.pt', 'dataset/raw_expert_prob': 0.5, 'dataset/raw_flip_prob_bounds': [0.01, 0.1], 'dataset/raw_git_sha': "unknown (Command '['git', '-C', '/root/simdist/UWLab', 'rev-parse', 'HEAD']' returned non-zero exit status 128.)", 'dataset/raw_image_encoding': 'JPEG (RGB); decode with PIL, or cv2.imdecode(...)[:, :, ::-1]', 'dataset/raw_insertive_object': 'JigV2', 'dataset/raw_jpeg_quality': 90, 'dataset/raw_noise_interval': [1, 5], 'dataset/raw_num_envs': 64, 'dataset/raw_num_rows': 32768, 'dataset/raw_proprio_dims': [1, 6, 6, 6, 1], 'dataset/raw_proprio_terms': ['last_gripper_action', 'last_arm_action', 'arm_joint_pos', 'end_effector_pose', 'binary_contact'], 'dataset/raw_receptive_object': 'BottomEnclosure', 'dataset/raw_row_order': 'step-major: num_envs consecutive rows per env step', 'dataset/raw_seed': 4242, 'dataset/raw_shard_index': 0, 'dataset/raw_sigma_bounds': [0.05, 1.0], 'dataset/raw_suboptimal_checkpoints': ['model_0.pt', 'model_100.pt', 'model_200.pt', 'model_300.pt', 'model_400.pt', 'model_500.pt', 'model_600.pt', 'model_700.pt', 'model_800.pt', 'model_900.pt', 'model_1000.pt', 'model_1100.pt', 'model_1200.pt', 'model_1300.pt', 'model_1400.pt', 'model_1500.pt', 'model_1600.pt', 'model_1700.pt', 'model_1800.pt', 'model_1900.pt', 'model_2000.pt', 'model_2100.pt', 'model_2200.pt', 'model_2300.pt', 'model_2400.pt', 'model_2500.pt', 'model_2600.pt', 'model_2700.pt', 'model_2800.pt', 'model_2900.pt', 'model_3000.pt', 'model_3100.pt', 'model_3200.pt', 'model_3300.pt', 'model_3400.pt', 'model_3500.pt', 'model_3600.pt'], 'dataset/raw_system': 'omnireset_ur10e', 'dataset/raw_task': 'OmniReset-UR10eLinearGripper-RelCartesianOSC-RGB-DataCollection-v0', 'dataset/raw_num_shards': 10, 'code/simdist_git_sha': None, 'code/simdist_git_dirty': None, 'instance/hostname': 'ebc50e8fee3f', 'instance/platform': 'Linux-6.8.0-90-generic-x86_64-with-glibc2.39', 'instance/python': '3.10.20', 'instance/cpu_count': 32, 'instance/cuda_visible_devices': '', 'instance/gpus': ['NVIDIA GeForce RTX 5090']}
+Parameters to be optimized: 52231753
+Total trainable parameters: 52231753
+Total parameters: 52241413
+Gradient clipping ON at global norm 1.0.
+Starting training
+Starting epoch 1/2
+
+  self.pid = os.fork()
+2026-08-13 04:19:24.280170: W external/xla/xla/tsl/framework/bfc_allocator.cc:310] Allocator (GPU_0_bfc) ran out of memory trying to allocate 9.51GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
+2026-08-13 04:19:25.827681: W external/xla/xla/tsl/framework/bfc_allocator.cc:310] Allocator (GPU_0_bfc) ran out of memory trying to allocate 20.01GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
+2026-08-13 04:19:31.419125: W external/xla/xla/tsl/framework/bfc_allocator.cc:310] Allocator (GPU_0_bfc) ran out of memory trying to allocate 33.55GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
+2026-08-13 04:19:35.582474: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng0{} for conv %cudnn-conv-bw-filter.20 = (f32[64,7,7,3]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.123167, %bitcast.123172), window={size=7x7 stride=2x2 pad=3_3x3_3}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convBackwardFilter", metadata={op_name="jit(train_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=767}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:19:39.452619: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 4.870214035s
+Trying algorithm eng0{} for conv %cudnn-conv-bw-filter.20 = (f32[64,7,7,3]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.123167, %bitcast.123172), window={size=7x7 stride=2x2 pad=3_3x3_3}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convBackwardFilter", metadata={op_name="jit(train_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=767}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+
+
+  self.pid = os.fork()
+
+Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/diagnostics/step_200
+[2026-08-13 04:21:39,355][absl][INFO] - [thread=MainThread] Failed to get flag value for EXPERIMENTAL_ORBAX_USE_DISTRIBUTED_PROCESS_ID.
+[2026-08-13 04:21:39,355][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
+[2026-08-13 04:21:39,355][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7795781073d0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7795781073d0>}).
+[2026-08-13 04:21:39,355][absl][INFO] - orbax-checkpoint version: 0.11.25
+[2026-08-13 04:21:39,356][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512.
+[2026-08-13 04:21:39,356][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/model_config.yaml) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:21:39,356][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/diagnostics) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:21:39,359][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/model_config.yaml) to be a directory.
+[2026-08-13 04:21:39,359][absl][INFO] - Found 0 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512
+[2026-08-13 04:21:39,360][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=None, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x77976500d0f0>
+[2026-08-13 04:21:39,360][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 04:21:39,360][absl][INFO] - [process=0] Saving checkpoint at step 200
+[2026-08-13 04:21:39,360][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.
+[2026-08-13 04:21:39,360][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp
+[2026-08-13 04:21:39,361][absl][INFO] - Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786594899361154757, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786594899361154757, "commit_timestamp_nsecs": null, "custom_metadata": {}} to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 04:21:39,361][absl][INFO] - Using ThreadSafeKeyValueSignalingClient
+[2026-08-13 04:21:39,361][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:21:39,362][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 04:21:39,362][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 04:21:39,362][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 04:21:39,363][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7796d3d729e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 04:21:39,363][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x779578107a60>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 04:21:39,363][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 04:21:39,363][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7796d3d729e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 04:21:39,363][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7795781075e0>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 04:21:39,387][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 04:21:39,441][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.054824s
+[2026-08-13 04:21:39,462][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 04:21:39,537][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 04:21:39,668][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.206075s
+[2026-08-13 04:21:39,673][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 199.3 MiB) (time elapsed: 309 milliseconds) (per-host)
+[2026-08-13 04:21:39,679][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.315151s (batch_requests_ready=0.016641s, total_serialization_initiated=0.288125s, others=0.010385s)
+[2026-08-13 04:21:39,681][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1 Bytes/s (total gbytes: 398.5 MiB) (time elapsed: 310 milliseconds) (per-host)
+[2026-08-13 04:21:39,681][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.311037s (batch_requests_ready=0.080589s, total_serialization_initiated=0.227885s, others=0.002564s)
+[2026-08-13 04:21:39,681][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.320059s (all_items=0.000311s, per_item={'default': '0.00022411', 'opt_state': '0.00008726'}, temp_paths=0.319747)
+[2026-08-13 04:21:39,681][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:21:40,026][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.344699s (commit=0.008488s, array_metadata_write=0.336211s)
+[2026-08-13 04:21:40,028][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 199.3 MiB) (time elapsed: 664 milliseconds) (per-host)
+[2026-08-13 04:21:40,029][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:21:40,096][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 04:21:40,238][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.557090s (commit=0.538428s, array_metadata_write=0.018662s)
+[2026-08-13 04:21:40,239][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 398.5 MiB) (time elapsed: 868 milliseconds) (per-host)
+[2026-08-13 04:21:40,240][absl][INFO] - Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786594899361154757, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 04:21:40,241][absl][INFO] - Updated Metadata={'item_handlers': {'default': 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler', 'opt_state': 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786594899361154757, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 04:21:40,243][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 04:21:40,246][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.004962s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 04:21:40,247][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/default
+[2026-08-13 04:21:40,249][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 04:21:40,253][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.005998s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 04:21:40,254][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp/opt_state
+[2026-08-13 04:21:40,254][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200
+[2026-08-13 04:21:40,255][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200`.
+[2026-08-13 04:21:40,255][absl][INFO] - Finished synchronous save in 0.90 seconds to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200
+[2026-08-13 04:21:40,255][absl][INFO] - [process=0][thread=MainThread][step=200] CheckpointManager Save Finalize is syncing with other hosts...
+[2026-08-13 04:21:40,256][absl][INFO] - [process=0][thread=MainThread][step=200] CheckpointManager Save Finalize is done on all hosts.
+[2026-08-13 04:21:40,256][absl][INFO] - [process=0][thread=MainThread][step=200] Finished synchronous save.
+[2026-08-13 04:21:40,256][absl][INFO] - {'step': 200, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786594899.3602803, 'wait_for_prev_duration_secs': 3.552436828613281e-05, 'checkpointer_blocking_start_time': 1786594899.3603978, 'checkpointer_blocking_duration_secs': 0.8954441547393799, 'get_old_steps_start_time': 1786594900.2558548, 'get_old_steps_duration_secs': 7.343292236328125e-05, 'checkpoint_manager_blocking_start_time': 1786594899.3602335, 'checkpoint_manager_blocking_duration_secs': 0.8958446979522705}
+[2026-08-13 04:21:40,256][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 04:21:40,256][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=131493256062784 count=1 at 0x7797655d3c80>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7795781059f0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x779578106890>, _write_futures=[])
+
+
+
+  self.pid = os.fork()
+
+Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/diagnostics/step_250
+[2026-08-13 04:22:06,415][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
+[2026-08-13 04:22:06,416][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7793623c0f40>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7793623c0f40>}).
+[2026-08-13 04:22:06,416][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512.
+[2026-08-13 04:22:06,417][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/model_config.yaml) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:22:06,417][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/diagnostics) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:22:06,417][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/200) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:22:06,419][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/model_config.yaml) to be a directory.
+[2026-08-13 04:22:06,421][absl][INFO] - Missing metrics for step 200
+[2026-08-13 04:22:06,421][absl][INFO] - Found 1 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512
+[2026-08-13 04:22:06,421][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=None, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x779337f7c5b0>
+[2026-08-13 04:22:06,421][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 04:22:06,422][absl][INFO] - [process=0] Saving checkpoint at step 250
+[2026-08-13 04:22:06,422][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.
+[2026-08-13 04:22:06,422][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp
+[2026-08-13 04:22:06,423][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:22:06,423][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 04:22:06,423][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 04:22:06,424][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 04:22:06,424][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7796d3d729e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 04:22:06,424][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7797651967a0>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 04:22:06,424][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 04:22:06,424][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7796d3d729e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 04:22:06,424][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7797651968f0>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 04:22:06,448][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 04:22:06,464][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.015794s
+[2026-08-13 04:22:06,479][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 04:22:06,521][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 04:22:06,563][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.084392s
+[2026-08-13 04:22:06,564][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1 Bytes/s (total gbytes: 199.3 MiB) (time elapsed: 138 milliseconds) (per-host)
+[2026-08-13 04:22:06,564][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.138185s (batch_requests_ready=0.016724s, total_serialization_initiated=0.121207s, others=0.000254s)
+[2026-08-13 04:22:06,569][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 2 Bytes/s (total gbytes: 398.5 MiB) (time elapsed: 137 milliseconds) (per-host)
+[2026-08-13 04:22:06,577][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.144990s (batch_requests_ready=0.031974s, total_serialization_initiated=0.100011s, others=0.013005s)
+[2026-08-13 04:22:06,581][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.158793s (all_items=0.000323s, per_item={'default': '0.00023031', 'opt_state': '0.00009227'}, temp_paths=0.158471)
+[2026-08-13 04:22:06,583][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:22:06,697][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 04:22:07,058][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.493958s (commit=0.166845s, array_metadata_write=0.327113s)
+[2026-08-13 04:22:07,060][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 199.3 MiB) (time elapsed: 634 milliseconds) (per-host)
+[2026-08-13 04:22:07,061][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:22:07,156][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.578494s (commit=0.560052s, array_metadata_write=0.018441s)
+[2026-08-13 04:22:07,157][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 398.5 MiB) (time elapsed: 724 milliseconds) (per-host)
+[2026-08-13 04:22:07,161][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 04:22:07,165][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.005843s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 04:22:07,165][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/default
+[2026-08-13 04:22:07,168][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 04:22:07,173][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.006834s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 04:22:07,174][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp/opt_state
+[2026-08-13 04:22:07,174][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250
+[2026-08-13 04:22:07,175][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250`.
+[2026-08-13 04:22:07,175][absl][INFO] - Finished synchronous save in 0.75 seconds to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512/250
+[2026-08-13 04:22:07,176][absl][INFO] - [process=0][thread=MainThread][step=250] CheckpointManager Save Finalize is syncing with other hosts...
+[2026-08-13 04:22:07,176][absl][INFO] - [process=0][thread=MainThread][step=250] CheckpointManager Save Finalize is done on all hosts.
+[2026-08-13 04:22:07,176][absl][INFO] - [process=0][thread=MainThread][step=250] Finished synchronous save.
+[2026-08-13 04:22:07,176][absl][INFO] - {'step': 250, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786594926.4219565, 'wait_for_prev_duration_secs': 3.8623809814453125e-05, 'checkpointer_blocking_start_time': 1786594926.4220767, 'checkpointer_blocking_duration_secs': 0.753901481628418, 'get_old_steps_start_time': 1786594927.1759915, 'get_old_steps_duration_secs': 7.2479248046875e-05, 'checkpoint_manager_blocking_start_time': 1786594926.4219046, 'checkpoint_manager_blocking_duration_secs': 0.7543132305145264}
+[2026-08-13 04:22:07,176][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 04:22:07,176][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=131493256062784 count=1 at 0x77976ab258c0>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7793623c09a0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7793623c0c40>, _write_futures=[])
+
+Steps: 250, Metrics: {'steps_per_second': 2.3066735478356386, 'throughput/samples_per_s': 1181.016856491847, 'lr': 4.980000085197389e-06, 'grad_norm': 14.506815738677979, 'dataloader_wait_frac': 0.004342438876728982, 'gpu_util': 100.0, 'gpu_mem_used_frac': 0.9340939062164566, 'train/loss': 2.1081204414367676, 'train/latent_dynamics': 0.6832031607627869, 'train/reward': 0.2987804114818573, 'train/value': 0.1729889214038849, 'train/action': 0.9531482458114624, 'test/loss': nan, 'test/latent_dynamics': nan, 'test/reward': nan, 'test/value': nan, 'test/action': nan, 'test/dataloader_wait_frac': 0.0, 'train/unweighted/latent_dynamics': 0.6832031607627869, 'train/unweighted/reward': 0.2987804114818573, 'train/unweighted/value': 0.1729889214038849, 'train/unweighted/action': 0.2382870614528656, 'test/unweighted/latent_dynamics': nan, 'test/unweighted/reward': nan, 'test/unweighted/value': nan, 'test/unweighted/action': nan, 'eval/decoder_probe/proprio_mse': nan, 'eval/decoder_probe/image_mse': nan, 'eval/value_auc_success': 0.7415237426757812, 'eval/value_auc_num_episodes': 512.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 1.0, 'epoch': 0}
+Stopping training after 250 steps.

demos/dimensionality_ladder/smoke_d512_r2.log

@@ -0,0 +1,205 @@
+/root/simdist/simdist_uwlab/scripts/train_model.py:14: UserWarning: 
+The version_base parameter is not specified.
+Please specify a compatability version level, or None.
+Will assume defaults for version 1.1
+  @hydra.main(**get_train_model_hydra_config())
+/root/simdist/venv_train/lib/python3.10/site-packages/hydra/_internal/hydra.py:119: UserWarning: Future Hydra versions will no longer change working directory at job runtime by default.
+See https://hydra.cc/docs/1.2/upgrades/1.1_to_1.2/changes_to_job_working_dir/ for more information.
+  ret = run_job(
+INFO:2026-08-13 04:37:37,037:jax._src.xla_bridge:867: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
+Training will stop after 250 steps.
+Running training with config:  {'system': {'name': 'omnireset_ur10e', 'proprio_obs': {'types': [{'name': 'last_gripper_action', 'dim': 1, 'noise': 0.01}, {'name': 'last_arm_action', 'dim': 6, 'noise': 0.01}, {'name': 'arm_joint_pos', 'dim': 6, 'noise': 0.005}, {'name': 'end_effector_pose', 'dim': 6, 'noise': 0.003}, {'name': 'binary_contact', 'dim': 1, 'noise': 0.05}]}, 'extero_obs': {'types': [{'name': 'front_rgb', 'dim': [120, 160, 3], 'noise': 0.01}, {'name': 'side_rgb', 'dim': [120, 160, 3], 'noise': 0.01}, {'name': 'wrist_rgb', 'dim': [120, 160, 3], 'noise': 0.01}]}, 'cmd': {'dim': 0}, 'actions': [{'name': 'ee_delta_x', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_y', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_z', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_rx', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_ry', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'ee_delta_rz', 'min_noise': 0.05, 'max_noise': 1.0}, {'name': 'gripper_binary', 'min_noise': 0.0, 'max_noise': 0.0, 'flip_prob_min': 0.01, 'flip_prob_max': 0.1}]}, 'model': {'type': 'manipulator_world_model', 'dataset': {'type': 'manipulator_world_model', 'history_length': 5, 'prediction_length': 5, 'jpeg_backend': 'auto', 'augmentations': {'add_noise': True, 'image': {'enabled': True, 'random_crop': {'enabled': True, 'padding': [4, 6]}, 'gaussian_blur': {'enabled': True, 'prob': 0.5, 'kernel_size': 5, 'sigma': [0.1, 1.5]}, 'color_jitter': {'enabled': True, 'brightness': 0.2, 'contrast': 0.2, 'saturation': 0.2, 'hue': 0.02}}}}, 'dropout': {'mlp': 0.2, 'attention': 0.05}, 'latent_dim': 512, 'encoder': {'mlp_hidden_size_factor': 2, 'proprio_obs_layers': 2, 'extero_obs': {'image_resnet': {'share_trunk': True, 'pretrained': True, 'freeze_batch_norm': True, 'normalize_input': True}}, 'action_layers': 0, 'latent_layers': 2, 'latent_norm': True}, 'embedding': {'mlp_hidden_size_factor': 2, 'future_acts_layers': 0, 'future_cmds_layers': 0, 'reward_layers': 0, 'value_layers': 0}, 'dynamics': {'attention': {'layers': 3, 'heads': 4, 'mlp_hidden_size_factor': 4, 'mask': 'causal'}}, 'reward': {'attention': {'layers': 1, 'heads': 1, 'mlp_hidden_size_factor': 4, 'mask': None}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}, 'value': {'attention': {'layers': 1, 'heads': 1, 'mlp_hidden_size_factor': 4, 'mask': None}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}, 'policy': {'attention': {'layers': 4, 'heads': 8, 'mlp_hidden_size_factor': 4, 'mask': 'causal'}, 'decoder': {'layers': 0, 'mlp_hidden_size_factor': 2}}}, 'loss': {'type': 'world_model', 'weights': {'latent_dynamics': 1.0, 'reward': 1.0, 'value': 1.0, 'action': 4.0}}, 'heads': {}, 'run_name': 'smoke_d512_r2', 'data': {'dataset_name': 'simdist_nodr', 'num_train_workers': 8, 'num_test_workers': 4, 'prefetch_factor': 2, 'paper_scale_steps': 25000000}, 'training': {'num_epochs': 100, 'max_steps': 250, 'eval_interval': 200, 'learning_rate': 0.0002, 'warmup_steps': 10000, 'decay_steps': 400000, 'end_learning_rate': 0.0001, 'batch_size': 512, 'training_data_ratio': 0.98, 'seed': 42, 'grad_clip_norm': 1.0}, 'checkpoint': {'enabled': True, 'max_to_keep': None, 'resume_checkpoint': None}, 'eval': {'max_batches': 8, 'scatter_points': 5000, 'episodes': {'enabled': True, 'holdout': 128, 'num_auc_episodes': 512, 'num_curve_episodes': 16, 'match_frac': 0.5, 'curve_stride': 5, 'batch_size': 32, 'success': {'metric': 'episode_return', 'threshold': None}}, 'diagnostics': {'enabled': True, 'dirname': 'diagnostics'}, 'decoder_probe': {'enabled': True, 'hidden_size': 128, 'learning_rate': 0.001, 'steps_per_eval': 200, 'image': {'enabled': True, 'base_channels': 128, 'learning_rate': 0.001, 'batch': 32}}}, 'wandb': {'log': False, 'project': 'simdist_train_model', 'entity': None, 'mode': 'online', 'dir': None}}
+Reserved 17360 chunks (128 whole episodes) for the episode probe; they are in neither the train nor the test split.
+Creating a new model...
+[2026-08-13 04:37:37,037][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
+Run fingerprint: {'dataset/processed_dir': '/root/simdist/simdist_uwlab/datasets/sim/simdist_nodr/processed_data_omnireset_ur10e_H-5_T-5', 'dataset/num_steps': 299985, 'dataset/num_trajectories': 261585, 'dataset/generation_git_sha': "unknown (Command '['git', '-C', '/root/simdist/UWLab', 'rev-parse', 'HEAD']' returned non-zero exit status 128.)", 'dataset/expert_checkpoint': 'model_3700.pt', 'dataset/paper_scale_steps': 25000000, 'dataset/paper_scale_fraction': 0.0119994, 'dataset/raw_camera_names': ['front_rgb', 'side_rgb', 'wrist_rgb'], 'dataset/raw_camera_resolution': [120, 160], 'dataset/raw_clean_interval': [5, 10], 'dataset/raw_created_utc': '2026-08-12T20:28:07.508201+00:00', 'dataset/raw_expert_checkpoint': 'model_3700.pt', 'dataset/raw_expert_prob': 0.5, 'dataset/raw_flip_prob_bounds': [0.01, 0.1], 'dataset/raw_git_sha': "unknown (Command '['git', '-C', '/root/simdist/UWLab', 'rev-parse', 'HEAD']' returned non-zero exit status 128.)", 'dataset/raw_image_encoding': 'JPEG (RGB); decode with PIL, or cv2.imdecode(...)[:, :, ::-1]', 'dataset/raw_insertive_object': 'JigV2', 'dataset/raw_jpeg_quality': 90, 'dataset/raw_noise_interval': [1, 5], 'dataset/raw_num_envs': 64, 'dataset/raw_num_rows': 32768, 'dataset/raw_proprio_dims': [1, 6, 6, 6, 1], 'dataset/raw_proprio_terms': ['last_gripper_action', 'last_arm_action', 'arm_joint_pos', 'end_effector_pose', 'binary_contact'], 'dataset/raw_receptive_object': 'BottomEnclosure', 'dataset/raw_row_order': 'step-major: num_envs consecutive rows per env step', 'dataset/raw_seed': 4242, 'dataset/raw_shard_index': 0, 'dataset/raw_sigma_bounds': [0.05, 1.0], 'dataset/raw_suboptimal_checkpoints': ['model_0.pt', 'model_100.pt', 'model_200.pt', 'model_300.pt', 'model_400.pt', 'model_500.pt', 'model_600.pt', 'model_700.pt', 'model_800.pt', 'model_900.pt', 'model_1000.pt', 'model_1100.pt', 'model_1200.pt', 'model_1300.pt', 'model_1400.pt', 'model_1500.pt', 'model_1600.pt', 'model_1700.pt', 'model_1800.pt', 'model_1900.pt', 'model_2000.pt', 'model_2100.pt', 'model_2200.pt', 'model_2300.pt', 'model_2400.pt', 'model_2500.pt', 'model_2600.pt', 'model_2700.pt', 'model_2800.pt', 'model_2900.pt', 'model_3000.pt', 'model_3100.pt', 'model_3200.pt', 'model_3300.pt', 'model_3400.pt', 'model_3500.pt', 'model_3600.pt'], 'dataset/raw_system': 'omnireset_ur10e', 'dataset/raw_task': 'OmniReset-UR10eLinearGripper-RelCartesianOSC-RGB-DataCollection-v0', 'dataset/raw_num_shards': 10, 'code/simdist_git_sha': None, 'code/simdist_git_dirty': None, 'instance/hostname': 'ebc50e8fee3f', 'instance/platform': 'Linux-6.8.0-90-generic-x86_64-with-glibc2.39', 'instance/python': '3.10.20', 'instance/cpu_count': 32, 'instance/cuda_visible_devices': '', 'instance/gpus': ['NVIDIA GeForce RTX 5090']}
+Parameters to be optimized: 52231753
+Total trainable parameters: 52231753
+Total parameters: 52241413
+Gradient clipping ON at global norm 1.0.
+Starting training
+Starting epoch 1/100
+
+  self.pid = os.fork()
+
+
+  self.pid = os.fork()
+2026-08-13 04:39:38.982000: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng1{k2=2,k3=0} for conv %cudnn-conv.59 = (f32[32,4,4,64]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6459, %bitcast.6469), window={size=59x79 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:39:39.112788: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.130810431s
+Trying algorithm eng1{k2=2,k3=0} for conv %cudnn-conv.59 = (f32[32,4,4,64]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6459, %bitcast.6469), window={size=59x79 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:39:42.211007: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng36{k2=0,k13=2,k14=3,k18=1,k23=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:39:42.447184: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.236240115s
+Trying algorithm eng36{k2=0,k13=2,k14=3,k18=1,k23=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:39:43.447259: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng28{k2=3,k3=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:39:44.191708: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.744494832s
+Trying algorithm eng28{k2=3,k3=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:39:45.191790: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng28{k2=4,k3=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:39:45.999878: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.808151788s
+Trying algorithm eng28{k2=4,k3=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:39:46.999971: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng1{k2=2,k3=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:40:01.849941: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 15.850031858s
+Trying algorithm eng1{k2=2,k3=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:40:02.850083: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng28{k2=1,k3=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+2026-08-13 04:40:04.727958: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 2.877947606s
+Trying algorithm eng28{k2=1,k3=0} for conv %cudnn-conv.60 = (f32[3,4,4,32]{3,2,1,0}, u8[0]{0}) custom-call(%bitcast.6604, %bitcast.6614), window={size=119x159 pad=1_1x1_1 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="jit(_image_probe_fit_step)/jit(main)/conv_general_dilated" source_file="/root/simdist/venv_train/lib/python3.10/site-packages/flax/nnx/nn/linear.py" source_line=989}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/diagnostics/step_200
+[2026-08-13 04:40:18,379][absl][INFO] - [thread=MainThread] Failed to get flag value for EXPERIMENTAL_ORBAX_USE_DISTRIBUTED_PROCESS_ID.
+[2026-08-13 04:40:18,379][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
+[2026-08-13 04:40:18,379][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7ca5e81cb070>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7ca5e81cb070>}).
+[2026-08-13 04:40:18,379][absl][INFO] - orbax-checkpoint version: 0.11.25
+[2026-08-13 04:40:18,380][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2.
+[2026-08-13 04:40:18,380][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:40:18,381][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/diagnostics) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:40:18,382][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/model_config.yaml) to be a directory.
+[2026-08-13 04:40:18,383][absl][INFO] - Found 0 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2
+[2026-08-13 04:40:18,383][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=None, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7ca5e82d6200>
+[2026-08-13 04:40:18,384][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 04:40:18,384][absl][INFO] - [process=0] Saving checkpoint at step 200
+[2026-08-13 04:40:18,384][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.
+[2026-08-13 04:40:18,384][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp
+[2026-08-13 04:40:18,384][absl][INFO] - Wrote Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786596018384765346, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786596018384765346, "commit_timestamp_nsecs": null, "custom_metadata": {}} to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 04:40:18,385][absl][INFO] - Using ThreadSafeKeyValueSignalingClient
+[2026-08-13 04:40:18,385][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:40:18,385][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 04:40:18,385][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 04:40:18,386][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 04:40:18,386][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7ca9cd57e9e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 04:40:18,386][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7ca5e81ca440>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 04:40:18,386][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 04:40:18,386][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7ca9cd57e9e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 04:40:18,386][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7ca5e81cb430>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 04:40:18,409][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 04:40:18,463][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.054024s
+[2026-08-13 04:40:18,481][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 04:40:18,550][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 04:40:18,665][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.183872s
+[2026-08-13 04:40:18,665][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 199.3 MiB) (time elapsed: 278 milliseconds) (per-host)
+[2026-08-13 04:40:18,677][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.290231s (batch_requests_ready=0.015528s, total_serialization_initiated=0.262689s, others=0.012013s)
+[2026-08-13 04:40:18,678][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 1 Bytes/s (total gbytes: 398.5 MiB) (time elapsed: 285 milliseconds) (per-host)
+[2026-08-13 04:40:18,679][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.286695s (batch_requests_ready=0.078769s, total_serialization_initiated=0.205727s, others=0.002198s)
+[2026-08-13 04:40:18,681][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.296154s (all_items=0.000284s, per_item={'default': '0.00019860', 'opt_state': '0.00008559'}, temp_paths=0.295870)
+[2026-08-13 04:40:18,687][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:40:18,715][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.044773s (commit=0.007656s, array_metadata_write=0.037117s)
+[2026-08-13 04:40:18,718][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 199.3 MiB) (time elapsed: 331 milliseconds) (per-host)
+[2026-08-13 04:40:18,718][absl][WARNING] - [process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:40:19,155][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 04:40:19,281][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.601949s (commit=0.583220s, array_metadata_write=0.018730s)
+[2026-08-13 04:40:19,282][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 398.5 MiB) (time elapsed: 889 milliseconds) (per-host)
+[2026-08-13 04:40:19,283][absl][INFO] - Read Metadata={'item_handlers': None, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786596018384765346, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} from /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 04:40:19,284][absl][INFO] - Updated Metadata={'item_handlers': {'default': 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler', 'opt_state': 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786596018384765346, 'commit_timestamp_nsecs': None, 'custom_metadata': {}} to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
+[2026-08-13 04:40:19,286][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 04:40:19,290][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.006000s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 04:40:19,290][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/default
+[2026-08-13 04:40:19,292][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 04:40:19,297][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.005991s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 04:40:19,297][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp/opt_state
+[2026-08-13 04:40:19,298][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200
+[2026-08-13 04:40:19,298][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200`.
+[2026-08-13 04:40:19,299][absl][INFO] - Finished synchronous save in 0.91 seconds to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200
+[2026-08-13 04:40:19,299][absl][INFO] - [process=0][thread=MainThread][step=200] CheckpointManager Save Finalize is syncing with other hosts...
+[2026-08-13 04:40:19,299][absl][INFO] - [process=0][thread=MainThread][step=200] CheckpointManager Save Finalize is done on all hosts.
+[2026-08-13 04:40:19,299][absl][INFO] - [process=0][thread=MainThread][step=200] Finished synchronous save.
+[2026-08-13 04:40:19,299][absl][INFO] - {'step': 200, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786596018.3841789, 'wait_for_prev_duration_secs': 3.695487976074219e-05, 'checkpointer_blocking_start_time': 1786596018.384294, 'checkpointer_blocking_duration_secs': 0.9151127338409424, 'get_old_steps_start_time': 1786596019.2994237, 'get_old_steps_duration_secs': 9.298324584960938e-05, 'checkpoint_manager_blocking_start_time': 1786596018.3841274, 'checkpoint_manager_blocking_duration_secs': 0.9156544208526611}
+[2026-08-13 04:40:19,299][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 04:40:19,299][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=137072306804544 count=1 at 0x7ca5f0363480>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7ca5e8156f20>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7ca5e81ca950>, _write_futures=[])
+
+
+
+  self.pid = os.fork()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/diagnostics/step_250
+[2026-08-13 04:40:53,170][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
+[2026-08-13 04:40:53,171][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7ca5b845d2d0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7ca5b845d2d0>}).
+[2026-08-13 04:40:53,171][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2.
+[2026-08-13 04:40:53,172][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:40:53,172][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/diagnostics) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:40:53,172][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/200) to end with ".orbax-checkpoint-tmp".
+[2026-08-13 04:40:53,174][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/model_config.yaml) to be a directory.
+[2026-08-13 04:40:53,176][absl][INFO] - Missing metrics for step 200
+[2026-08-13 04:40:53,176][absl][INFO] - Found 1 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2
+[2026-08-13 04:40:53,176][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=None, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7ca5b845c4c0>
+[2026-08-13 04:40:53,176][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 04:40:53,176][absl][INFO] - [process=0] Saving checkpoint at step 250
+[2026-08-13 04:40:53,176][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.
+[2026-08-13 04:40:53,177][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp
+[2026-08-13 04:40:53,178][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:40:53,178][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 04:40:53,178][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 04:40:53,178][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 04:40:53,178][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7ca9cd57e9e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 04:40:53,178][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7ca5b845e560>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 04:40:53,178][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
+[2026-08-13 04:40:53,179][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7ca9cd57e9e0>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
+[2026-08-13 04:40:53,179][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7ca5b845e050>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
+[2026-08-13 04:40:54,312][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 04:40:54,328][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.016510s
+[2026-08-13 04:40:54,338][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
+[2026-08-13 04:40:54,389][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 04:40:54,432][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.094213s
+[2026-08-13 04:40:54,446][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 199.3 MiB) (time elapsed: a second) (per-host)
+[2026-08-13 04:40:54,449][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.269997s (batch_requests_ready=0.016129s, total_serialization_initiated=1.246996s, others=0.006872s)
+[2026-08-13 04:40:54,462][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 398.5 MiB) (time elapsed: a second) (per-host)
+[2026-08-13 04:40:54,464][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.278361s (batch_requests_ready=1.142724s, total_serialization_initiated=0.121144s, others=0.014494s)
+[2026-08-13 04:40:54,466][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 1.288386s (all_items=0.000303s, per_item={'default': '0.00021291', 'opt_state': '0.00009036'}, temp_paths=1.288083)
+[2026-08-13 04:40:54,466][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:40:54,553][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
+[2026-08-13 04:40:54,610][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.160592s (commit=0.148358s, array_metadata_write=0.012233s)
+[2026-08-13 04:40:54,611][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 199.3 MiB) (time elapsed: a second) (per-host)
+[2026-08-13 04:40:54,612][absl][WARNING] - [process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
+[2026-08-13 04:40:54,690][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.223852s (commit=0.208202s, array_metadata_write=0.015650s)
+[2026-08-13 04:40:54,691][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 398.5 MiB) (time elapsed: a second) (per-host)
+[2026-08-13 04:40:54,694][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 04:40:54,698][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.005563s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
+[2026-08-13 04:40:54,698][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/default
+[2026-08-13 04:40:54,700][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
+[2026-08-13 04:40:54,705][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.005764s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
+[2026-08-13 04:40:54,705][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp/opt_state
+[2026-08-13 04:40:54,706][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250
+[2026-08-13 04:40:54,707][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250`.
+[2026-08-13 04:40:54,707][absl][INFO] - Finished synchronous save in 1.53 seconds to /root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2/250
+[2026-08-13 04:40:54,707][absl][INFO] - [process=0][thread=MainThread][step=250] CheckpointManager Save Finalize is syncing with other hosts...
+[2026-08-13 04:40:54,708][absl][INFO] - [process=0][thread=MainThread][step=250] CheckpointManager Save Finalize is done on all hosts.
+[2026-08-13 04:40:54,708][absl][INFO] - [process=0][thread=MainThread][step=250] Finished synchronous save.
+[2026-08-13 04:40:54,708][absl][INFO] - {'step': 250, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/smoke_d512_r2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786596053.1766784, 'wait_for_prev_duration_secs': 3.814697265625e-05, 'checkpointer_blocking_start_time': 1786596053.1767955, 'checkpointer_blocking_duration_secs': 1.531048059463501, 'get_old_steps_start_time': 1786596054.707862, 'get_old_steps_duration_secs': 7.462501525878906e-05, 'checkpoint_manager_blocking_start_time': 1786596053.176628, 'checkpoint_manager_blocking_duration_secs': 1.5314555168151855}
+[2026-08-13 04:40:54,708][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
+[2026-08-13 04:40:54,708][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=137072306804544 count=1 at 0x7ca6143c28c0>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7ca5b845cdc0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7ca5b845ce50>, _write_futures=[])
+
+Steps: 250, Metrics: {'steps_per_second': 2.323871326697501, 'throughput/samples_per_s': 1189.8221192691205, 'lr': 4.980000085197389e-06, 'grad_norm': 21.387936582565306, 'dataloader_wait_frac': 0.004600914697782326, 'gpu_util': 100.0, 'gpu_mem_used_frac': 0.933909896647959, 'train/loss': 3.1581077575683594, 'train/latent_dynamics': 0.7485272288322449, 'train/reward': 0.4558477997779846, 'train/value': 0.7240086793899536, 'train/action': 1.2297242879867554, 'test/loss': 1.9590742588043213, 'test/latent_dynamics': 0.5563856363296509, 'test/reward': 0.2872110903263092, 'test/value': 0.16700537502765656, 'test/action': 0.9484719038009644, 'test/dataloader_wait_frac': 0.1465656357537205, 'train/unweighted/latent_dynamics': 0.7485272288322449, 'train/unweighted/reward': 0.4558477997779846, 'train/unweighted/value': 0.7240086793899536, 'train/unweighted/action': 0.30743107199668884, 'test/unweighted/latent_dynamics': 0.5563856363296509, 'test/unweighted/reward': 0.2872110903263092, 'test/unweighted/value': 0.16700537502765656, 'test/unweighted/action': 0.2371179759502411, 'eval/latent_rollout_error/h1': 0.5557060688734055, 'eval/latent_rollout_error/h2': 0.5499007999897003, 'eval/latent_rollout_error/h3': 0.5607725009322166, 'eval/latent_rollout_error/h4': 0.5593321099877357, 'eval/latent_rollout_error/h5': 0.5533562526106834, 'eval/latent_rollout_error_mean': 0.5558135464787484, 'eval/latent_variance/encoder': 0.08250186778604984, 'eval/latent_variance/predicted': 0.12474203761667013, 'eval/reward_r2': 0.5703360972413711, 'eval/reward_explained_variance': 0.6051282913693438, 'eval/reward_pearson_r': 0.7779296503761246, 'eval/value_r2': 0.62845755069178, 'eval/value_explained_variance': 0.6792964932054467, 'eval/value_pearson_r': 0.8242125486001747, 'eval/bc_action_mse': 0.9307147652596321, 'eval/num_batches': 8.0, 'eval/decoder_probe/proprio_mse': 0.6365486979484558, 'eval/decoder_probe/image_mse': 0.01991439051926136, 'eval/value_auc_success': 0.72900390625, 'eval/value_auc_num_episodes': 128.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 1.0, 'epoch': 0}
+Stopping training after 250 steps.