(bsize)
| 123 | """ |
| 124 | |
| 125 | def try_batch_size(bsize): |
| 126 | print(f"Attempting batch size: {bsize}") |
| 127 | # Spawn a worker for each attempt |
| 128 | attempt_q = Queue() |
| 129 | p = Process( |
| 130 | target=test_batch_size_worker, |
| 131 | args=(attempt_q, model_name, input_ids, attention_mask, bsize, device, use_xformers) |
| 132 | ) |
| 133 | p.start() |
| 134 | p.join() |
| 135 | result = attempt_q.get() |
| 136 | p = None |
| 137 | if result[0] == 'error': |
| 138 | # If there's an error unrelated to OOM, raise it |
| 139 | print(f"Error occurred: {result[1]}") |
| 140 | raise RuntimeError(result[1]) |
| 141 | success = result[1] |
| 142 | print(f"Batch size {bsize}: {'succeeded' if success else 'failed'}") |
| 143 | |
| 144 | print("Clearing CUDA cache and garbage collection") |
| 145 | torch.cuda.empty_cache() |
| 146 | gc.collect() |
| 147 | return success |
| 148 | |
| 149 | try: |
| 150 | print("\nStarting batch size search...") |
no outgoing calls
no test coverage detected