Add ReFT (LoReFT, NoReFT, DiReFT)

[calpt]

This PR integrates multiple ReFT variants as new adapter methods.

Paper: https://arxiv.org/pdf/2404.03592
Original code: https://github.com/stanfordnlp/pyreft

Changes

• Add ReFT module implementation via ReftLayer, integrated into all models supported by Adapters. Integration via init_reft() method & Pytorch hook.
• Add new ReftConfig as base config class with three default instances: LoReftConfig, NoReftConfig and DiReftConfig.
• Method documentation can be found here: https://github.com/adapter-hub/adapters/blob/6c19ea06c143621a735226e477bf772068e55be3/docs/methods.md#reft

Compatibility

Tested that Pyreft & Adapters produce the same outputs on inference by converting Pyreft checkpoints to Adapters checkpoints (tested settings: LoReft, NoReft, DiReft, weight tying, prefix, suffix, rank, mostly using roberta-base).

Script for testing & checkpoint conversion here: https://github.com/calpt/pyreft/blob/main/compatibility.py.

Evaluation

Roberta-base with LoReFT on GLUE, using hyperparameters similar to the paper:

Task	Score
Cola (Matthews Corr.)	53.95
MNLI (Acc.)	83.23
MRPC (F1)	91.70
QNLI (Acc.)	90.94
QQP (Acc.)	86.82
RTE (Acc.)	76.53
SST-2 (Acc.)	93.81
STS-B (Spearmanr)	88.99

Todos

• Modeling implementations
• Add test methods
• Make all checks passing
• Add documentation
• Make sure implementation produces same outputs as original code
• Sanity check training runs

[frankaging]

@calpt Thanks for the PR! I took a quick look, and it looks promising. Here are two minor questions:

1. It seems like tied and untied weights are handled here:
   https://github.com/adapter-hub/adapters/pull/705/files#diff-e791dd9c62ff127d32170821d7571b69e641f7ecc9462a58227fa5a80f3502f1R80

Does this mean if I want united weights among prefix and suffix, I will create two adaptors and set self.prefix_positions or self.suffix_positions to none?

2. Since in the ReFT paper, we only run experiments where we add interventions to the residual stream (e.g., transformer layer or block output), will, by default, ReftConfig assume this? If I want to add interventions to different streams (e.g., attention output, MLP up-projection), how will the ReftConfig look like?

Thanks!

[calpt]

@frankaging Thanks for looking over this! Re your questions:

Re 1: In the current implementation, one or two modules per layer will be created depending on whether the tied_weights attribute in the config is set or not. See here:

adapters/src/adapters/methods/reft.py

Lines 50 to 62 in 25f1d1c

	n_units = 1 if config.tied_weights else 2
	self.units = nn.ModuleList(
	[
	ReftUnit(
	in_features,
	config.r,
	config.orthogonality,
	config.subtract_projection,
	config.non_linearity,
	config.dropout,
	)
	for _ in range(n_units)
	]

This makes it very easy for a user to to tie or not tie weights when adding a single Reft adapter, e.g.:

from adapters import AutoAdapterModel, ReftConfig

model = AutoAdapterModel.from_pretrained("...")

config = ReftConfig(
    layers="all", prefix_positions=1, suffix_positions=1, r=1,
    tied_weights=True  # set to True or False to share weights
)
model.add_adapter("my_reft", config=config)
model.set_active_adapters("my_reft")

Re 2: Currently, the Reft implementation always assumes interventions are added to the residual stream as you explained, since this is the method proposed in the paper. This is done via a PyTorch hook here:

adapters/src/adapters/methods/reft.py

Lines 162 to 169 in 25f1d1c

	def init_reft(model):
	def hook_fn(module, args, output):
	return (module.reft_layer(output[0]),) + output[1:]
	for _, layer in model.iter_layers():
	if not hasattr(layer, "reft_layer"):
	layer.reft_layer = ReftLayer(model.config, model.adapters_config)
	layer.register_forward_hook(hook_fn)

While no other intervention points are added for now, we can easily extend with similar hooks for other intervention points where it makes sense to do so.

Thanks again for looking over this! Please let us know if you have any suggestions or ideas what we should add or change for the first version!
Currently, the PR is still in a draft state. Once remaining issues are fixed and we have some documentation, would be happy to get your feedback again before we merge.

[frankaging]

@calpt Thanks for your responses! It makes sense to me.

Will the hook work out of the box for accelerated training (e.g., deepspeed, etc..)? Any existing tests on this? Thanks!

[calpt]

@frankaging No extensive tests for training yet at this point. Deepspeed support of this library is unfortunately flaky in general and not really a focus at the moment, but using e.g. torch distributed or HF Accelerate should work in the end.

[hSterz]

This looks great! Just some small comments and questions

[hSterz]

Do we want to leave this TODO open?

[calpt]

added, pls check if this is correct

[lenglaender]

Commented on some minor things; everything else looks good & correctly implemented to me

Once the open comments are resolved & left padding is implemented this is ready to merge.

[lenglaender]

Looks good to me. This is ready to merge

[calpt]

Thanks! I've added some quick training results on GLUE tasks in the description, based on that the implementation looks good.

@frankaging re distributed training: I've tested it to work with torch distributed & HF Accelerate via the Trainer class, e.g. for GLUE:

torchrun --standalone --nnodes=1 --nproc-per-node=2 examples/pytorch/text-classification/run_glue.py \
  --model_name_or_path roberta-large \
  --task_name $TASK_NAME \
  --do_train \
  --do_eval \
  --max_seq_length $SEQ \
  --pad_to_max_length False \
  --per_device_train_batch_size 32 \
  --learning_rate $LR \
  --warmup_ratio $WARMUP \
  --num_train_epochs $EPOCH \
  --output_dir output/$TASK_NAME \
  --overwrite_output_dir \
  --train_adapter \
  --adapter_config "loreft[prefix_positions=$POS]"

[frankaging]

Thanks! I've added some quick training results on GLUE tasks in the description, based on that the implementation looks good.

@frankaging re distributed training: I've tested it to work with torch distributed & HF Accelerate via the Trainer class, e.g. for GLUE:

torchrun --standalone --nnodes=1 --nproc-per-node=2 examples/pytorch/text-classification/run_glue.py \
  --model_name_or_path roberta-large \
  --task_name $TASK_NAME \
  --do_train \
  --do_eval \
  --max_seq_length $SEQ \
  --pad_to_max_length False \
  --per_device_train_batch_size 32 \
  --learning_rate $LR \
  --warmup_ratio $WARMUP \
  --num_train_epochs $EPOCH \
  --output_dir output/$TASK_NAME \
  --overwrite_output_dir \
  --train_adapter \
  --adapter_config "loreft[prefix_positions=$POS]"

@calpt Thanks! It's great to see this approach works for different kinds of parallel training!

I was re-looking into that orthogonal matrix initialization thing (i.e., I was referencing the PEFT repo ticket and asking whether we should remove a redundant init), and I found that for some cases, removing that init step might cause unstable results. Have you look in to this again by doing some tests on your side? Thanks.

[calpt]

I was re-looking into that orthogonal matrix initialization thing (i.e., I was referencing the PEFT repo ticket and asking whether we should remove a redundant init), and I found that for some cases, removing that init step might cause unstable results. Have you look in to this again by doing some tests on your side? Thanks.

Interesting, I haven't tested this specifically. From looking at the code, it makes sense the orthogonal init is redundant, would you suggest we re-add it still?