Skip to content

Commit

Permalink
remove prints for now
Browse files Browse the repository at this point in the history
  • Loading branch information
@ahmeda14960
ahmeda14960 committed Feb 1, 2025
1 parent 05fbc29 commit 75bcaee
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 34 deletions.
42 changes: 21 additions & 21 deletions src/levanter/data/packing.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,57 +113,57 @@ def pack_prompt_completions(
Packs a list of prompt completions into LmExamples using the SequencePacker
"""

start_packer = time.perf_counter()
# start_packer = time.perf_counter()
packers = [SequencePacker(Pos, max_segments_per_example, pad_token)]
print(f"time to create packer is {time.perf_counter() - start_packer:.6f} seconds", flush=True)
# print(f"time to create packer is {time.perf_counter() - start_packer:.6f} seconds", flush=True)
# put timer in here
for sequence in sequences:
start_time = time.perf_counter()
# start_time = time.perf_counter()
# put timer here, keep in mind the loop is an iterator so we want to
# separate the time to get examples, and how long it takes to pack
loss_mask = np.arange(len(sequence.ids)) >= sequence.prompt_length - 1
loss_mask[-1] = 0
assert np.any(loss_mask)

# time how long to pack, and subtract
start_for_pack_yield = time.perf_counter()
#start_for_pack_yield = time.perf_counter()
for packer in packers:
if packer.can_pack(sequence.ids):
add_example = time.perf_counter()
# add_example = time.perf_counter()
packer.add_example(sequence.ids, loss_mask, sequence.segment_id)
add_example_end = time.perf_counter()
time_to_add_example = add_example_end - add_example
print(f" time to add example to segment packer { time_to_add_example:.6f} seconds", flush=True)
end_iter_plus_example = time.perf_counter()
time_to_reach_example_total = end_iter_plus_example - start_time
print(f"time to get sequence, so time to get here minus time to add example {time_to_reach_example_total - time_to_add_example:.6f} seconds", flush=True)
# add_example_end = time.perf_counter()
# time_to_add_example = add_example_end - add_example
# print(f" time to add example to segment packer { time_to_add_example:.6f} seconds", flush=True)
# end_iter_plus_example = time.perf_counter()
# time_to_reach_example_total = end_iter_plus_example - start_time
# print(f"time to get sequence, so time to get here minus time to add example {time_to_reach_example_total - time_to_add_example:.6f} seconds", flush=True)
if packer.num_segments == max_segments_per_example:
ot = packer.pack()
end_time = time.perf_counter()
time_to_yield = end_time - start_for_pack_yield
print(f"MAX SEG total time to for loop until we yielded a packed example {time_to_yield}", flush=True)
print(f"MAX SEG total time for iterator {start_time - time_to_yield}", flush=True)
# end_time = time.perf_counter()
# time_to_yield = end_time - start_for_pack_yield
# print(f"MAX SEG total time to for loop until we yielded a packed example {time_to_yield}", flush=True)
# print(f"MAX SEG total time for iterator {start_time - time_to_yield}", flush=True)
yield ot
packers.remove(packer)
break
else:
# no packer could fit the example, create a new one
start_new_packer = time.perf_counter()
#start_new_packer = time.perf_counter()
packer = SequencePacker(Pos, max_segments_per_example, pad_token)
packer.add_example(sequence.ids, loss_mask, sequence.segment_id)
packers.append(packer)
print(f"time to create new packer is {time.perf_counter() - start_new_packer:.6f}", flush=True)
#print(f"time to create new packer is {time.perf_counter() - start_new_packer:.6f}", flush=True)

while len(packers) >= max_buffered_examples:
start_return_packed_example_full_buffer = time.perf_counter()
#start_return_packed_example_full_buffer = time.perf_counter()
max_example = packers.pop(0).pack()
print(f"time to create lm example when max buffered examples is {time.perf_counter() - start_return_packed_example_full_buffer:.6f}", flush=True)
#print(f"time to create lm example when max buffered examples is {time.perf_counter() - start_return_packed_example_full_buffer:.6f}", flush=True)
yield max_example

for packer in packers:
start_return_packed_example = time.perf_counter()
#start_return_packed_example = time.perf_counter()
example = packer.pack()
print(f"time to create lm example from packer is {time.perf_counter() - start_return_packed_example:.6f}", flush=True)
#print(f"time to create lm example from packer is {time.perf_counter() - start_return_packed_example:.6f}", flush=True)
yield example


Expand Down
26 changes: 13 additions & 13 deletions src/levanter/main/sft.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -288,11 +288,11 @@ def create_prompt_completion_iterator(cached_dataset: AsyncDataset, Pos: hax.Axi
# TODO play around with batch size
for batch_indicies in batched(range(length), 4096):
# put timer here
start_time = time.perf_counter()
# start_time = time.perf_counter()
examples = asyncio.run(cached_dataset.get_batch(batch_indicies))
end_time = time.perf_counter()
elapsed_time = end_time - start_time
print(f"Elapsed time for get batches: {elapsed_time:.6f} seconds", flush=True)
# end_time = time.perf_counter()
# elapsed_time = end_time - start_time
# print(f"Elapsed time for get batches: {elapsed_time:.6f} seconds", flush=True)

for i in range(len(examples)):
example = examples[i]
Expand Down Expand Up @@ -357,21 +357,21 @@ def stack_batches(example_iterator, Pos, TrainBatch):
batch_count = 0
for batch in batched(example_iterator, TrainBatch.size):
batch_count += 1
start_time_loop = time.perf_counter()
#start_time_loop = time.perf_counter()
if len(batch) < TrainBatch.size:
dummy_instance = _make_dummy_instance(batch, Pos)
batch.extend([dummy_instance] * (TrainBatch.size - len(batch)))
# Start timing before calling stack_tree
start_time = time.perf_counter()
# # Start timing before calling stack_tree
# start_time = time.perf_counter()
result = stack_tree(TrainBatch, batch) # Capture the result
stack_time = time.perf_counter() - start_time # Calculate elapsed time
# stack_time = time.perf_counter() - start_time # Calculate elapsed time

print(f"Stack tree execution time: {stack_time:.6f} seconds", flush=True)
# print(f"Stack tree execution time: {stack_time:.6f} seconds", flush=True)
yield result # Yield the computed result
end_time_loop = time.perf_counter()
loop_time = end_time_loop - start_time_loop
print(f"Loop takes {loop_time}")
print(f"Iterator time is {loop_time - stack_time}")
# end_time_loop = time.perf_counter()
# loop_time = end_time_loop - start_time_loop
# print(f"Loop takes {loop_time}")
# print(f"Iterator time is {loop_time - stack_time}")

def add_special_tokens(tokenizer, use_unk_instead_of_adding=False):
special_tokens_dict = dict()
Expand Down

0 comments on commit 75bcaee

Please sign in to comment.