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.

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

write to S3 is very slow #812

Open
charliedream1 opened this issue Oct 25, 2024 · 3 comments
Open

write to S3 is very slow #812

charliedream1 opened this issue Oct 25, 2024 · 3 comments
Labels
bug Something isn't working

Comments

@charliedream1
Copy link

charliedream1 commented Oct 25, 2024

Environment

  • OS: [Ubuntu 20.04]
  • Hardware (GPU, or instance type): [H800]

To reproduce

I got a 2G jsonl.gz text file, I tokenized it with data stored as numpy array. Writer is defined as below:

   out = MDSWriter(
        columns={"input_ids": f"ndarray:int32:{args.seq_len}",
                 "token_type_ids": f"ndarray:int8:{args.seq_len}",
                 "attention_mask": f"ndarray:int8:{args.seq_len}",
                 "special_tokens_mask": f"ndarray:int8:{args.seq_len}"},
        out=out_path,
        compression=None
    )

tokenized data is pre-processed and loaded already, so time won't be wasted at tokenization.

def parse_data_2_mds_format(tokenized_dataset):
    input_ids = np.array(tokenized_dataset['input_ids']).astype(np.int32)
    token_type_ids = np.array(tokenized_dataset['token_type_ids']).astype(np.int8)
    attention_mask = np.array(tokenized_dataset['attention_mask']).astype(np.int8)
    special_tokens_mask = np.array(tokenized_dataset['special_tokens_mask']).astype(np.int8)
    return {'input_ids': input_ids, 'token_type_ids': token_type_ids,
            'attention_mask': attention_mask, 'special_tokens_mask': special_tokens_mask}

 with Pool(processes=args.mds_num_workers) as inner_pool:
                with tqdm(total=len(tokenized_datasets), desc="Writing Out MDS File") as pbar:
                    for result in inner_pool.imap(parse_data_2_mds_format, tokenized_datasets):
                        out.write(result)
                        pbar.update()

With code above, the writing time to S3 takes 30 minutes with mds_num_workers set as 200. If I set to 1, it
takes 1 hour to finish. It's just so slow, I have huge data to process. How to accelerate? Is that possible to write
a block of data once rather than one by one? Please give some suggestion to accelerate.

Expected behavior

Additional context

@charliedream1 charliedream1 added the bug Something isn't working label Oct 25, 2024
@snarayan21
Copy link
Collaborator

Hey @charliedream1, have you tried the parallel dataset conversion approach as detailed in our docs below? https://docs.mosaicml.com/projects/streaming/en/stable/preparing_datasets/parallel_dataset_conversion.html

Please let us know if that works for you.

@charliedream1
Copy link
Author

charliedream1 commented Oct 25, 2024 via email

@charliedream1
Copy link
Author

charliedream1 commented Oct 25, 2024 via email

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

No branches or pull requests

2 participants