fine-tune-dp runtime error in sample-level-dp of examples files

my running command:
python -m torch.distributed.run --nproc_per_node 4 examples/nlg-reddit/sample-level-dp/fine-tune-dp.py
--output_dir /u/nkp2mr/Songwei/dp-transformers-main/examples/nlg-reddit/sample-level-dp/finetune_private1
--model_name meta-llama/Llama-3.2-1B
--per_device_train_batch_size 32
--gradient_accumulation_steps 2
--evaluation_strategy steps
--eval_steps 200
--log_level info
--per_device_eval_batch_size 8
--eval_accumulation_steps 1
--seed 42
--target_epsilon 8
--per_sample_max_grad_norm 1.0
--prediction_loss_only
--weight_decay 0.01
--remove_unused_columns False
--num_train_epochs 30
--logging_steps 5
--max_grad_norm 0
--lr_scheduler_type constant
--learning_rate 1e-4
--disable_tqdm False
--dataloader_num_workers 2

error message:

57%|█████▋ | 500/870 [34:03<23:39, 3.84s/it][INFO|trainer.py:3503] 2024-10-03 15:12:07,734 >> Saving model checkpoint to /u/nkp2mr/Songwei/dp-transformers-main/examples/nlg-reddit/sample-level-dp/finetune_private1/checkpoint-500
[INFO|trainer.py:3517] 2024-10-03 15:12:07,736 >> Trainer.model is not a PreTrainedModel, only saving its state dict.

{'final_epsilon_prv': 5.837189086840157, 'final_epsilon_rdp': 6.5569718999375235, 'epoch': 17.24}

57%|█████▋ | 500/870 [34:03<23:39, 3.84s/it][rank0]: Traceback (most recent call last):
[rank0]: File "/u/nkp2mr/Songwei/dp-transformers-main/examples/nlg-reddit/sample-level-dp/fine-tune-dp_my.py", line 160, in
[rank0]: main(Arguments(train=train_args, privacy=privacy_args, model=model_args, lora=lora_args))
[rank0]: File "/u/nkp2mr/Songwei/dp-transformers-main/examples/nlg-reddit/sample-level-dp/fine-tune-dp_my.py", line 148, in main
[rank0]: trainer.train()
[rank0]: File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/transformers/trainer.py", line 1938, in train
[rank0]: return inner_training_loop(
[rank0]: ^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/transformers/trainer.py", line 2356, in _inner_training_loop
[rank0]: self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
[rank0]: File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/transformers/trainer.py", line 2807, in _maybe_log_save_evaluate
[rank0]: self._save_checkpoint(model, trial, metrics=metrics)
[rank0]: File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/transformers/trainer.py", line 2886, in _save_checkpoint
[rank0]: self.save_model(output_dir, _internal_call=True)
[rank0]: File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/transformers/trainer.py", line 3454, in save_model
[rank0]: self._save(output_dir)
[rank0]: File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/transformers/trainer.py", line 3519, in _save
[rank0]: safetensors.torch.save_file(
[rank0]: File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/safetensors/torch.py", line 286, in save_file
[rank0]: serialize_file(_flatten(tensors), filename, metadata=metadata)
[rank0]: ^^^^^^^^^^^^^^^^^
[rank0]: File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/safetensors/torch.py", line 488, in _flatten
[rank0]: raise RuntimeError(
[rank0]: RuntimeError:
[rank0]: Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: [{'_module.module.base_model.model.lm_head.weight', '_module.module.base_model.model.model.embed_tokens.weight'}].
[rank0]: A potential way to correctly save your model is to use save_model.
[rank0]: More information at https://huggingface.co/docs/safetensors/torch_shared_tensors
[rank0]:

57%|█████▋ | 500/870 [34:03<25:12, 4.09s/it]
[rank0]:[W1003 15:12:08.638436808 ProcessGroupNCCL.cpp:1168] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present, but this warning has only been added since PyTorch 2.4 (function operator())
W1003 15:12:12.355000 139967302018240 torch/distributed/elastic/multiprocessing/api.py:858] Sending process 3186093 closing signal SIGTERM
W1003 15:12:12.355000 139967302018240 torch/distributed/elastic/multiprocessing/api.py:858] Sending process 3186094 closing signal SIGTERM
W1003 15:12:12.356000 139967302018240 torch/distributed/elastic/multiprocessing/api.py:858] Sending process 3186095 closing signal SIGTERM
E1003 15:12:14.625000 139967302018240 torch/distributed/elastic/multiprocessing/api.py:833] failed (exitcode: 1) local_rank: 0 (pid: 3186092) of binary: /u/nkp2mr/anaconda3/envs/dp-transformers/bin/python
Traceback (most recent call last):
File "", line 198, in _run_module_as_main
File "", line 88, in _run_code
File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/torch/distributed/run.py", line 905, in
main()
File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/init.py", line 348, in wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/torch/distributed/run.py", line 901, in main
run(args)
File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/torch/distributed/run.py", line 892, in run
elastic_launch(
File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 133, in call
return launch_agent(self._config, self._entrypoint, list(args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/u/nkp2mr/anaconda3/envs/dp-transformers/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 264, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:

examples/nlg-reddit/sample-level-dp/fine-tune-dp_my.py FAILED

Failures:
<NO_OTHER_FAILURES>

Root Cause (first observed failure):
[0]:
time : 2024-10-03_15:12:12
host : dplab08
rank : 0 (local_rank: 0)
exitcode : 1 (pid: 3186092)
error_file: <N/A>
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html

This error always happens in 500 steps, and the check point file always has only one, which is "check-point-500". It seems like this is a problem about saving model.

Start with examples/nlg-reddit/sample-level-dp/fine-tune-dp.py and compare it with the fine-tune-dp_my.py file named in the traceback. Reproduce the distributed command far enough to inspect checkpoint saving at step 500 and trace the safetensors shared-tensor error. Done means the run can save a usable checkpoint and continue past that step.