Issue of Trainer on WMDP

[chengjiali]

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task
• My own task or dataset (give details below)

Reproduction

Hi, thank you for adding the WMDP dataset.

The recent changes are merged and the dataset is setup. I tried to run an experiment on WMDP with the following command.

#!/bin/bash

export MASTER_PORT=$(python -c "import socket; s=socket.socket(); s.bind(('', 0)); print(s.getsockname()[1]); s.close()")
echo "Master Port: $MASTER_PORT"

per_device_train_batch_size=4
gradient_accumulation_steps=2

model=zephyr-7b-beta

data_splits=(
    "cyber"
    "bio"
)

trainers=(
    "GradAscent"
    "GradDiff"
    "NPO"
    "SimNPO"
)


for data_split in "${data_splits[@]}"; do
    for trainer in "${trainers[@]}"; do

            task_name=wmdp_${model}_${data_split}_${trainer}

            CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch --config_file configs/accelerate/default_config.yaml --main_process_port $MASTER_PORT \
                src/train.py --config-name=unlearn.yaml \
                experiment=unlearn/wmdp/default.yaml \
                model=${model} \
                data_split=${data_split} \
                trainer=${trainer} \
                task_name=${task_name} \
                paths.output_dir=saves/unlearn/${task_name}/ \
                trainer.args.per_device_train_batch_size=${per_device_train_batch_size} \
                trainer.args.gradient_accumulation_steps=${gradient_accumulation_steps} \
                trainer.args.ddp_find_unused_parameters=true \
                trainer.args.gradient_checkpointing=true
    done
done

Error:
For GA

  File "/user/open-unlearning/src/train.py", line 48, in main
    trainer, trainer_args = load_trainer(
                            ^^^^^^^^^^^^^
  File "/user/open-unlearning/src/trainer/__init__.py", line 64, in load_trainer
    trainer = trainer_cls(
              ^^^^^^^^^^^^
  File "/user/open-unlearning/src/trainer/base.py", line 19, in __init__
    super().__init__(*args, **kwargs)
TypeError: Trainer.__init__() got an unexpected keyword argument 'gamma'

For GD

Traceback (most recent call last):
  File "/user/open-unlearning/src/train.py", line 48, in main
    trainer, trainer_args = load_trainer(
                            ^^^^^^^^^^^^^
  File "/user/open-unlearning/src/trainer/__init__.py", line 64, in load_trainer
    trainer = trainer_cls(
              ^^^^^^^^^^^^
  File "/user/open-unlearning/src/trainer/unlearn/grad_diff.py", line 8, in __init__
    super().__init__(*args, **kwargs)
  File "/user/open-unlearning/src/trainer/base.py", line 19, in __init__
    super().__init__(*args, **kwargs)
TypeError: Trainer.__init__() got an unexpected keyword argument 'steering_coeff'

Could you help me look into this? Could you please provide an official script for WMDP unlearning? Thank you so much for your help.

Expected behavior

Run with no bug.

[c0syfeng]

same issue！ Could you share your official scripts for WMDP unlearning? thanks

[Dornavineeth]

Hi, The file unlearn/wmdp/default.yaml currently sets RMU as the default method, along with related parameters like gamma, steering_coeff, etc. If you override the method in your bash script using trainer=GradAsc, these RMU-specific parameters will still be set and raise the error, since they are not compatible with GradAsc. Better to change the default method for WMDP to GradAsc.

If you wish to work on your bash script, you can make the below modifications in unlearn/wmdp/default.yaml

defaults:
  - override /trainer: GradAsc

# remove the trainer.method_args

Please feel free to raise a PR, If you get a chance to make this change and test it.

[chengjiali]

Thank you! using the updated unlearn/wmdp/default.yaml fixes this issue.

I met anther problem. When doing unlearning on WMDP-Bio, I noticed a huge memory consumption in data processing, upto 1000G RAM.

Is there anything to mitigate the RAM consumption? Thanks.

[c0syfeng]

@Dornavineeth hi, thanks for your reply! But, when I run scripts followed your advice, I encountered this bug:

(unlearning) (base) root@autodl-container-27b944a7b2-13cdc01c:~/autodl-tmp/open-unlearning# bash scripts/wmdp_unlearn.sh
Master Port: 46975
[2025-06-09 11:23:41,063] [INFO] [real_accelerator.py:219:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-09 11:23:45,207] [INFO] [real_accelerator.py:219:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-09 11:23:45,219] [INFO] [real_accelerator.py:219:get_accelerator] Setting ds_accelerator to cuda (auto detect)
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
Loading checkpoint shards: 100%|███████████████████████████████████████████████████| 8/8 [00:01<00:00,  5.49it/s]
Loading checkpoint shards: 100%|███████████████████████████████████████████████████| 8/8 [00:01<00:00,  5.51it/s]
Generating train split: 60887 examples [00:07, 8578.46 examples/s]
W0609 12:13:31.841000 140664732002112 torch/distributed/elastic/multiprocessing/api.py:858] Sending process 8419 closing signal SIGTERM
E0609 12:13:36.662000 140664732002112 torch/distributed/elastic/multiprocessing/api.py:833] failed (exitcode: -9) local_rank: 0 (pid: 8418) of binary: /root/miniconda3/envs/unlearning/bin/python3.11
Traceback (most recent call last):
  File "/root/miniconda3/envs/unlearning/bin/accelerate", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/root/miniconda3/envs/unlearning/lib/python3.11/site-packages/accelerate/commands/accelerate_cli.py", line 48, in main
    args.func(args)
  File "/root/miniconda3/envs/unlearning/lib/python3.11/site-packages/accelerate/commands/launch.py", line 1159, in launch_command
    deepspeed_launcher(args)
  File "/root/miniconda3/envs/unlearning/lib/python3.11/site-packages/accelerate/commands/launch.py", line 852, in deepspeed_launcher
    distrib_run.run(args)
  File "/root/miniconda3/envs/unlearning/lib/python3.11/site-packages/torch/distributed/run.py", line 892, in run
    elastic_launch(
  File "/root/miniconda3/envs/unlearning/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 133, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/unlearning/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 264, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
=====================================================
src/train.py FAILED
-----------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
-----------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2025-06-09_12:13:31
  host      : autodl-container-27b944a7b2-13cdc01c
  rank      : 0 (local_rank: 0)
  exitcode  : -9 (pid: 8418)
  error_file: <N/A>
  traceback : Signal 9 (SIGKILL) received by PID 8418
=====================================================

Could you help me fix that? that looks like OOM error, but my gpu wasn't used during training.

my Nvidia-smi:

Every 2.0s: nvidia-smi                             autodl-container-27b944a7b2-13cdc01c: Mon Jun  9 12:31:34 2025

Mon Jun  9 12:31:34 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.14              Driver Version: 550.54.14      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA A800 80GB PCIe          On  |   00000000:4F:00.0 Off |                  Off |
| N/A   32C    P0             44W /  300W |       0MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA A800 80GB PCIe          On  |   00000000:D1:00.0 Off |                  Off |
| N/A   36C    P0             46W /  300W |       0MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

packages of my env are as follows:

(unlearning) (base) root@autodl-container-27b944a7b2-13cdc01c:~/autodl-tmp/open-unlearning# pip list
Package                  Version     Build
------------------------ ----------- -----
absl-py                  2.3.0
accelerate               0.34.2
aiohappyeyeballs         2.6.1
aiohttp                  3.12.11
aiosignal                1.3.2
annotated-types          0.7.0
antlr4-python3-runtime   4.9.3
attrs                    25.3.0
bitsandbytes             0.44.1
certifi                  2025.4.26
chardet                  5.2.0
charset-normalizer       3.4.2
click                    8.2.1
colorama                 0.4.6
colorlog                 6.9.0
DataProperty             1.1.0
datasets                 3.0.1
deepspeed                0.15.4
dill                     0.3.8
einops                   0.8.1
evaluate                 0.4.3
filelock                 3.18.0
flash-attn               2.6.3
frozenlist               1.6.2
fsspec                   2024.6.1
grpcio                   1.72.1
hjson                    3.1.0
huggingface-hub          0.29.1
hydra-colorlog           1.2.0
hydra-core               1.3.0
idna                     3.10
Jinja2                   3.1.6
joblib                   1.5.1
jsonlines                4.0.0
lm_eval                  0.4.8
lxml                     5.4.0
Markdown                 3.8
MarkupSafe               3.0.2
mbstrdecoder             1.1.4
more-itertools           10.7.0
mpmath                   1.3.0
msgpack                  1.1.0
multidict                6.4.4
multiprocess             0.70.16
networkx                 3.5
ninja                    1.11.1.4
nltk                     3.9.1
numexpr                  2.10.2
numpy                    2.2.3
nvidia-cublas-cu12       12.1.3.1
nvidia-cuda-cupti-cu12   12.1.105
nvidia-cuda-nvrtc-cu12   12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12        9.1.0.70
nvidia-cufft-cu12        11.0.2.54
nvidia-curand-cu12       10.3.2.106
nvidia-cusolver-cu12     11.4.5.107
nvidia-cusparse-cu12     12.1.0.106
nvidia-nccl-cu12         2.20.5
nvidia-nvjitlink-cu12    12.9.86
nvidia-nvtx-cu12         12.1.105
omegaconf                2.3.0
open-unlearning          0.1.0
packaging                25.0
pandas                   2.3.0
pathvalidate             3.2.3
peft                     0.15.2
pip                      25.1
portalocker              3.1.1
propcache                0.3.1
protobuf                 6.31.1
psutil                   7.0.0
py-cpuinfo               9.0.0
pyarrow                  20.0.0
pybind11                 2.13.6
pydantic                 2.11.5
pydantic_core            2.33.2
pytablewriter            1.2.1
python-dateutil          2.9.0.post0
pytz                     2025.2
PyYAML                   6.0.2
regex                    2024.11.6
requests                 2.32.3
rouge_score              0.1.2
sacrebleu                2.5.1
safetensors              0.5.3
scikit-learn             1.5.2
scipy                    1.14.1
setuptools               78.1.1
six                      1.17.0
sqlitedict               2.1.0
sympy                    1.14.0
tabledata                1.3.4
tabulate                 0.9.0
tcolorpy                 0.1.7
tensorboard              2.18.0
tensorboard-data-server  0.7.2
threadpoolctl            3.6.0
tokenizers               0.20.3
torch                    2.4.1
tqdm                     4.67.1
tqdm-multiprocess        0.0.11
transformers             4.45.1
triton                   3.0.0       1
typepy                   1.3.4
typing_extensions        4.14.0
typing-inspection        0.4.1
tzdata                   2025.2
urllib3                  2.4.0
Werkzeug                 3.1.3
wheel                    0.45.1
word2number              1.1
xxhash                   3.5.0
yarl                     1.20.0
zstandard                0.23.0

[Dornavineeth]

@chengjiali Yes. The current pretraining dataset implementation is a bit inefficient and might consume a lot of CPU memory. Not sure if it is 1000GB, may I know how you calculated the memory consumption?
But yes, there is a scope for better implementation. Do let us know if you have any thoughts for implementation

@c0syfeng, Was it happening while loading the dataset or while training? Could it be the CPU OOM for processing the WMDP dataset?

[chengjiali]

Hi, thanks for your reply.

I just eyeball it on htop on a machine with 1000G RAM. Sometimes the job will be killed due to OOM, but a bit random.

Does using datasets.Dataset.map in place of _chunk_text make this more efficient? I guess at least it caches the processed data for faster load.

[c0syfeng]

@Dornavineeth it happening while loading pretrain dataset. After debugging, I found the OOM occured at this
full_token_sequence = self.tokenizer(raw_text, add_special_tokens=False)[ "input_ids" ]

raw_text is too huge which leads to collapse. Are there any approaches to address this issues?

[ywglf]

@Dornavineeth Besides, because of the pad_token is the same as the eos_token in zephyr-7b-beta tokennizer, though the padding_side = left, the flash attention will still raise an error, the attention_mask generated from input_ids should be modified in src/data/collators.py line49.
"You are attempting to perform batched generation with padding_side='right'"
" this may lead to unexpected behaviour for Flash Attention version of Mistral. Make sure to "
" call tokenizer.padding_side = 'left' before tokenizing the input. "

[Dornavineeth]

@ywglf I couldn't understand the issue. Can you please explain what you tried, and what errors were raised?
What change are you suggesting in line49?

If its related to padding_side while tokenization, you can change it configs/experiment/unlearn/wmdp/default.yaml.

[Dornavineeth]

@c0syfeng I think it's just the CPU memory issue. One way to handle large datasets is by using IterableDataset. I haven't tested the following code yet, but you can give it a try. Let me know if it works.

from torch.utils.data import IterableDataset

class PretrainingDataset(IterableDataset):
    def __init__(self, hf_args, template_args, tokenizer, text_key="text", max_length=2048):
        super(PretrainingDataset, self).__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.text_key = text_key
        self.dataset = load_dataset(**hf_args, streaming=True)[hf_args.get("split", "train")]

    def __iter__(self):
        buffer = []
        buffer_len = 0

        for example in self.dataset:
            text = example[self.text_key]
            input_ids = self.tokenizer.encode(text, add_special_tokens=False)
            buffer.extend(input_ids)
            buffer_len += len(input_ids)

            while buffer_len >= self.max_length:
                chunk = buffer[:self.max_length]
                buffer = buffer[self.max_length:]
                buffer_len -= self.max_length

                chunk_text = self.tokenizer.decode(chunk)
                yield preprocess_pretraining_instance(
                    self.tokenizer, "", chunk_text, self.max_length
                )

[ywglf]

@ywglf I couldn't understand the issue. Can you please explain what you tried, and what errors were raised? What change are you suggesting in line49?

If its related to padding_side while tokenization, you can change it configs/experiment/unlearn/wmdp/default.yaml.

When the EOS token and PAD token are the same, attention mask=input_ids. ne (self. tokenizer. pad_token. id) will also mark the EOS token at the end of the sequence as 0 (padding), which is a serious issue for left padding. Flash Attention will misunderstand the padding side.

Add code below among src/data/collators.py line49 and line50 can temporarily alleviate the problem.
if self.padding_side == "left":
attention_mask[:, -1] = True

[c0syfeng]

@ywglf thx for your reminder! I also find the eos_token and pad_token of zephyr are both "</s>". When padding_size = 'left', the two line code you provided is necessary.

[chengjiali]

Hi @ywglf, thank you for providing this solution.

After adding these lines, the job is still killed due to OOM. Is there any other workaround?

Thank you.