Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update train.py to match structure of HuggingFaceH4/oasst1_en #113

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Update train.py to match structure of HuggingFaceH4/oasst1_en #113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
38d44c6
Update train.py
sarthak405 Aug 2, 2023
865d6d8
cache dir change
sarthak405 Oct 18, 2023
bfb6d54
Merge branch 'main' of https://github.com/sarthak405/starcoder
sarthak405 Oct 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion chat/deepspeed_z3_config_bf16.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_batch_size": 32,
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
28 changes: 15 additions & 13 deletions chat/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,13 +104,13 @@ def main():
###############
# Load datasets
###############
raw_datasets = load_dataset(data_args.dataset_name)
raw_datasets = load_dataset(data_args.dataset_name, cache_dir="/data-starcoder/cache")
logger.info(
f"Training on the following datasets and their proportions: {[split + ' : ' + str(dset.num_rows) for split, dset in raw_datasets.items()]}"
)
with training_args.main_process_first(desc="Log a few random samples from the raw training set"):
for index in random.sample(range(len(raw_datasets["train"])), 3):
logger.info(f"Sample {index} of the raw training set:\n\n{raw_datasets['train'][index]['messages']}")
for index in random.sample(range(len(raw_datasets["train_ift"])), 3):
logger.info(f"Sample {index} of the raw training set:\n\n{raw_datasets['train_ift'][index]['messages']}")

#########################
# Apply dialogue template
Expand All @@ -125,6 +125,7 @@ def main():
tokenizer = AutoTokenizer.from_pretrained(
model_args.model_name_or_path,
revision=model_args.model_revision,
cache_dir="/data-starcoder/cache"
)

# Note that we must call `add_tokens` before adding any special tokens
Expand All @@ -133,14 +134,14 @@ def main():
logger.info(f"Added {num_added_tokens} new tokens: {dialogue_tokens}")

if training_args.do_train:
column_names = list(raw_datasets["train"].features)
column_names = list(raw_datasets["train_ift"].features)
else:
column_names = list(raw_datasets["test"].features)
column_names = list(raw_datasets["test_ift"].features)
text_column_name = "text" if "text" in column_names else column_names[0]

with training_args.main_process_first(desc="Log a few random samples from the training set"):
for index in random.sample(range(len(raw_datasets["train"])), 3):
logger.info(f"Sample {index} of the raw training set:\n\n{raw_datasets['train'][index]['text']}")
for index in random.sample(range(len(raw_datasets["train_ift"])), 3):
logger.info(f"Sample {index} of the raw training set:\n\n{raw_datasets['train_ift'][index]['text']}")

# since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
Expand Down Expand Up @@ -218,17 +219,17 @@ def group_texts(examples):
)

if training_args.do_train:
if "train" not in tokenized_datasets:
if "train_ift" not in tokenized_datasets:
raise ValueError("--do_train requires a train dataset")
train_dataset = lm_datasets["train"]
train_dataset = lm_datasets["train_ift"]
if data_args.max_train_samples is not None:
max_train_samples = min(len(train_dataset), data_args.max_train_samples)
train_dataset = train_dataset.select(range(max_train_samples))

if training_args.do_eval:
if "test" not in tokenized_datasets:
if "test_ift" not in tokenized_datasets:
raise ValueError("--do_eval requires a validation dataset")
eval_dataset = lm_datasets["test"]
eval_dataset = lm_datasets["test_ift"]
if data_args.max_eval_samples is not None:
max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
eval_dataset = eval_dataset.select(range(max_eval_samples))
Expand All @@ -245,6 +246,7 @@ def group_texts(examples):
revision=model_args.model_revision,
torch_dtype=torch_dtype,
use_cache=False if training_args.gradient_checkpointing else True,
cache_dir="/data-starcoder/cache"
)
model.resize_token_embeddings(len(tokenizer))

Expand Down Expand Up @@ -281,8 +283,8 @@ def group_texts(examples):
)
metrics["train_samples"] = min(max_train_samples, len(train_dataset))

trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.log_metrics("train_ift", metrics)
trainer.save_metrics("train_ift", metrics)
trainer.save_state()

##########
Expand Down