From 041a6e31f25bcd3b40b9a9fc7101c1ca787b03ab Mon Sep 17 00:00:00 2001 From: Juhan Bae Date: Thu, 11 Jul 2024 11:01:14 -0400 Subject: [PATCH] Update commands --- examples/openwebtext/compute_scores.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/openwebtext/compute_scores.py b/examples/openwebtext/compute_scores.py index 53ec647..76d7d2d 100644 --- a/examples/openwebtext/compute_scores.py +++ b/examples/openwebtext/compute_scores.py @@ -92,8 +92,9 @@ def main(): analyzer.set_dataloader_kwargs(dataloader_kwargs) scores_name = args.factor_strategy + rank = args.query_gradient_rank if args.query_gradient_rank != -1 else None score_args = extreme_reduce_memory_score_arguments( - damping_factor=None, module_partitions=1, query_gradient_low_rank=args.query_gradient_rank, dtype=torch.bfloat16 + damping_factor=None, module_partitions=1, query_gradient_low_rank=rank, dtype=torch.bfloat16 ) score_args.module_partitions = 2 score_args.num_query_gradient_accumulations = 10