Implement batch inference for multi-apex

[EduardoTerres]

Use batch inference for efficiency

Edit:

when we perform inference for multi-apex we essentially do a for loop

    gpt_model, tokenizer = get_model_and_tokenizer(force_download, device=device)
    found_domains = set()

    for apex in sorted(domain_groups):
        found_domains |= _get_domains_for_group(...)

    return sorted(found_domains)

Here we are not making use of batch inference, as in we don't batch different apex related subdomains simultaneously.

[klaasmeinke]

batching is already implemented no?

https://github.com/hadriansecurity/subwiz/blob/main/subwiz/model.py#L340

[Martovark]

@klaasmeinke
Hi!
Perhaps it is suggested to add batch inference in this part:

    subs = [dom.subdomain for dom in input_domains]
    tokenizer_input = ",".join(sorted(subs)) + "[DELIM]"
    # TODO: pick a different subset, if some were out of context last iteration
    x = tokenizer.encode(tokenizer_input)
    x = [1] * (gpt_model.config.block_size - len(x)) + x
    x = torch.tensor(x)

It would be convenient to add different subdomains into single batch to speed up the search, but due to beam search, the most likely trajectories from one subdomain may displace another.

[EduardoTerres]

batching is already implemented no?

https://github.com/hadriansecurity/subwiz/blob/main/subwiz/model.py#L340

Indeed, the model can be trained and inferenced in batches.

However, when we perform inference for multi-apex we essentially do a for loop

    gpt_model, tokenizer = get_model_and_tokenizer(force_download, device=device)
    found_domains = set()

    for apex in sorted(domain_groups):
        found_domains |= _get_domains_for_group(...)

    return sorted(found_domains)

Here we are not making use of batch inference, as in we don't batch different apex related subdomains simultaneously.

[Martovark]

Thank you for the explanation.
Let me give you an example when it make sense to process different apex in single batch.

Suppose we have two different apexes, apex_1 and apex_2, after stacking:

gitlab,nginx[DELIM] // <- apex_1
github,apache[DELIM] // <- apex_2

I think it is justified to submit them together to the generate method if we potentially have millions of apexes and want to reduce processing time. I understand that in this situation, the most likely trajectories from apex_1 may displace apex_2 during beam search. But to speed things up, I'm willing to overlook that.

Plus in the case of large numbers of apex, if you don't want to save all generations of new subdomains after each step of generate, it is advisable to do DNS resolution in the same cycle to save only valid subdomains. This also takes time.

[klaasmeinke]

My intuition is that there not a lot of performance to be gained here because of the following back of the envelope calculations:

We use a default setting of n=500, which means that sequences_per_iter = round(5 + 500 / 4) = 126

From experimentation we have set batch_size = 8, because we found that this is roughly as time performant as using larger batch sizes (e.g. 32, 128) and more memory performant, allowing subwiz to be run on smaller hardware.

So the number of batches we inference per step are 15 batches of size 8 and 1 batch of 6 (using default settings, which I believe most user are using). So I would assume that the performance increase of running 16 batches of 8 is not a huge step up from running 15.75 batches of 8, and probably wouldn't warrant the extra code complexity. Happy to be proved wrong on this matter with metrics!

Regarding resolving before storing all subdomains per apex, I might have misunderstand your intent @Martovark, but it seems that we already do this on line 171 in the _get_domains_per_group function. Is that correct?

If you are suggesting that we do resolution in the .generate method itself between steps, then it is impossible to know what are the 500 most likely sequences before the .generate() method is finished running, so then we would be running resolution also for subdomains that are not in the 500 most likely subdomains. This would result in extra DNS resolutions and also would limit the parallelization of resolutions, since we resolve in smaller batches.

[Martovark]

I completely agree with you if standard settings (n=500, batch_size=8) are used for batch inference :)

I'll just describe my experiment to clarify the situation:

• About 10 million known domains, which I want to expand further.
• Removed the code with sequences_per_iter = round... and experimented with n=512, 1024, 2048, 8192 (n=8192 the maximum number, because size of the model vocab is 8192 and maximum categories 2^24 at which torch.multinomial does not drop).
• The model is quite small, so I set batch_size = topn in generate method.
• It seemed to me from a small number of tests that there is not much difference in the number of new domains between 1024, 2048, 8192, especially if we also take into account aliases. That's why I set n=1024.
• After every 1e6 new domains, I run dns resolve to keep only the new domains.

Overall, this process was not very fast.
So the first thing I noticed was:

x = [1] * (gpt_model.config.block_size - len(x)) + x

And I thought that in this place it is possible to speed up the search and better utilize cuda :)