실제 시나리오에서 모든 오픈 소스 LLM에 대한 강력한 트래픽 표현을 학습하고 다양한 트래픽 분석 작업 전반에 걸쳐 일반화를 향상시키는 범용 LLM 적응 프레임워크인 TrafficLLM 저장소입니다.
참고: 이 코드는 ChatGLM2 및 Llama2를 기반으로 합니다. 저자들에게 깊은 감사를 드립니다.
[2024.10.28] ChatGLM2보다 튜닝 및 추론 속도가 더 빠른 TrafficLLM을 구축하기 위해 GLM4를 사용하기 위한 적응 코드를 업데이트했습니다. 자세한 내용을 보려면 Adapt2GLM4로 이동하세요.
TrafficLLM은 자연어 및 트래픽 데이터를 사용하는 정교한 미세 조정 프레임워크를 기반으로 구축되었으며, 이는 네트워크 트래픽 분석에서 대규모 언어 모델의 유용성을 향상하기 위해 다음 기술을 제안합니다.
트래픽 도메인 토큰화. 자연어와 이종 트래픽 데이터 간의 양식 차이를 극복하기 위해 TrafficLLM은 트래픽 도메인 토큰화를 도입하여 LLM 적응을 위한 트래픽 감지 및 생성 작업의 다양한 입력을 처리합니다. 이 메커니즘은 대규모 트래픽 도메인 말뭉치에 대한 전문화된 토큰화 모델을 교육하여 LLM의 기본 토크나이저를 효과적으로 확장합니다.
듀얼 스테이지 튜닝 파이프라인. TrafficLLM은 다양한 트래픽 도메인 작업 전반에 걸쳐 LLM의 강력한 표현 학습을 달성하기 위해 이중 단계 튜닝 파이프라인을 사용합니다. 파이프라인은 다양한 트래픽 감지 및 생성 작업에 대한 TrafficLLM 작업 이해 및 트래픽 추론 능력을 기반으로 다양한 단계에서 지침을 이해하고 작업 관련 트래픽 패턴을 학습하도록 LLM을 교육합니다.
매개변수 효과적인 미세 조정(EA-PEFT)을 통한 확장 가능한 적응. 새로운 트래픽 환경에 대한 일반화를 위해 LLM을 적용하기 위해 TrafficLLM은 낮은 오버헤드로 모델 매개변수를 업데이트하는 매개변수 효과적인 미세 조정(EA-PEFT)을 사용한 확장 가능한 적응을 제안합니다. 이 기술은 모델 기능을 다양한 PEFT 모델로 분할하여 트래픽 패턴 변경으로 인해 발생하는 동적 시나리오의 비용을 최소화하는 데 도움이 됩니다.
우리는 다양한 교통 분석 작업 전반에 걸쳐 LLM 적응을 위한 400만 개 이상의 교통 데이터와 9,000개 이상의 인간 지침이 포함된 TrafficLLM의 교육 데이터 세트를 출시했습니다.
Instruction Datasets
: 지침 데이터 세트는 LLM이 트래픽 감지 또는 생성 작업에 대한 도메인 지식을 학습하고 다양한 시나리오에서 어떤 작업을 수행해야 하는지 이해하는 데 사용됩니다.
Traffic Datasets
: 교통 데이터 세트에는 공공 교통 데이터 세트에서 추출한 교통 튜닝 데이터가 포함되어 있어 LLM이 다양한 다운스트림 작업에서 트래픽 패턴을 학습하는 데 도움이 됩니다.
TrafficLLM에서 인간의 지시로 자연어 코퍼스를 구축하기 위해 전문가와 AI 보조자가 감독하는 9,209개의 작업별 지시를 수집했습니다. 통계는 다음과 같이 표시됩니다.
주류 업무 | 다운스트림 작업 | 약어 | #견본 |
---|---|---|---|
교통감지 | 악성코드 트래픽 탐지 | MTD | 1.0K |
봇넷 탐지 | BND | 1.1K | |
악성 DoH 탐지 | MDD | 0.6K | |
웹 공격 탐지 | 뭉치 | 0.6K | |
APT 공격 탐지 | AAD | 0.6K | |
암호화된 VPN 탐지 | EVD | 1.2K | |
Tor 행동 탐지 | 미정 | 0.6K | |
암호화된 앱 분류 | EAC | 0.6K | |
웹사이트 핑거프린팅 | WF | 0.6K | |
컨셉 드리프트 | CD | 0.6K | |
트래픽 생성 | 악성코드 트래픽 생성 | MTG | 0.6K |
봇넷 트래픽 생성 | BTG | 0.1K | |
암호화된 VPN 생성 | EVG | 0.4K | |
암호화된 앱 생성 | EAG | 0.6K |
다양한 네트워크 시나리오에서 TrafficLLM의 성능을 평가하기 위해 공개적으로 사용 가능한 트래픽 데이터 세트에서 040만 개가 넘는 튜닝 데이터를 추출하여 악성 트래픽과 양성 트래픽을 감지하거나 생성하는 TrafficLLM의 능력을 측정했습니다. 통계는 다음과 같이 표시됩니다.
데이터 세트 | 작업 | 약어 | #견본 |
---|---|---|---|
USTC TFC 2016 | 악성코드 트래픽 탐지 | MTD | 50.7K |
ISCX 봇넷 2014 | 봇넷 탐지 | BND | 25.0K |
DoHBrw 2020 | 악성 DoH 탐지 | MDD | 47.8K |
CSIC 2010 | 웹 공격 탐지 | 뭉치 | 34.5K |
DAPT 2020 | APT 공격 탐지 | AAD | 10.0K |
ISCX VPN 2016 | 암호화된 VPN 탐지 | EVD | 64.8K |
ISCX 토르 2016 | Tor 행동 탐지 | 미정 | 40.0K |
CSTNET 2023 | 암호화된 앱 분류 | EAC | 97.6K |
CW-100 2018 | 웹사이트 핑거프린팅 | WF | 7.4K |
APP-53 2023 | 컨셉 드리프트 | CD | 109.8K |
1. 환경 준비 2. TrafficLLM 교육 2.1. 사전 훈련된 체크포인트 준비 2.2. 데이터세트 전처리 2.3. 트래픽 도메인 토크나이저 훈련(선택 사항) 2.4. 신경언어 학습 튜닝 2.5. 작업별 트래픽 튜닝 2.6. PEFT(EA-PEFT)를 통한 확장 가능한 적응 3. TrafficLLM 평가 3.1. 체크포인트 및 데이터 준비 3.2. 평가 실행목차:
1. 환경 준비 [맨 위로 이동]
다음 명령을 실행하여 리포지토리를 복제하고 필요한 환경을 설치하세요.
conda create -n Trafficllm python=3.9 conda activate Trafficllm# TrafficLLMgit 복제본 복제 https://github.com/ZGC-LLM-Safety/TrafficLLM.gitcd TrafficLLM# 필요한 라이브러리 설치pip install -r 요구 사항.txt# trainingpip install rouge_chinese nltk jieba 데이터 세트
TrafficLLM은 세 가지 핵심 기술을 사용합니다. 즉, 지침 및 트래픽 데이터를 처리하는 트래픽 도메인 토큰화 , 텍스트 의미를 이해하고 다양한 작업에 걸쳐 트래픽 패턴을 학습하는 이중 단계 튜닝 파이프라인 , 업데이트하는 EA-PEFT 입니다. 새로운 시나리오 적응을 위한 모델 매개변수.
TrafficLLM은 기존 오픈 소스 LLM을 기반으로 훈련됩니다. 지침에 따라 체크포인트를 준비하세요.
ChatGLM2
: 간단한 배포 요구 사항을 갖춘 오픈 소스 LLM인 기본 모델 ChatGLM을 준비합니다. 여기에서 가중치를 다운로드하세요. 우리는 일반적으로 6B 매개변수를 갖는 v2 모델을 활용합니다.
Other LLMs
: 트래픽 분석 작업에 다른 LLM을 적용하려면 저장소의 교육 데이터를 재사용하고 공식 지침에 따라 교육 스크립트를 수정할 수 있습니다. 예를 들어 Llama2는 구성에 새 데이터 세트를 등록하는 데 필요합니다.
원시 트래픽 데이터 세트에서 LLM 학습에 적합한 교육 데이터를 추출하기 위해 다양한 작업에 대한 트래픽 데이터 세트를 전처리하는 특수 추출기를 설계합니다. 전처리 코드에는 구성에 대한 다음 매개변수가 포함되어 있습니다.
input
: 원시 트래픽 데이터 세트 경로(레이블이 지정된 하위 디렉터리가 포함된 기본 디렉터리 경로입니다. 레이블이 지정된 각 하위 디렉터리에는 전처리할 원시 .pcap 파일이 포함되어 있습니다).
dataset_name
: 원시 트래픽 데이터 세트 이름(이름이 TrafficLLM 코드에 등록되었는지 판단하는 데 도움이 됩니다).
traffic_task
: 탐지 작업 또는 생성 작업입니다.
granularity
: 패킷 수준 또는 흐름 수준 세분성입니다.
output_path
: 출력 훈련 데이터세트 경로입니다.
output_name
: 출력 훈련 데이터 세트 이름입니다.
패킷 수준 트래픽 탐지 작업을 위해 원시 트래픽 데이터 세트를 전처리하는 인스턴스입니다.
CD 전처리 python preprocess_dataset.py --input /Your/Raw/Dataset/Path --dataset_name /Your/Raw/Dataset/Name --traffic_task 감지 --세분성 패킷 --output_path /Your/Output/Dataset/Path --output_name /Your /Output/Dataset/Name
TrafficLLM은 신경 언어 및 데이터를 처리하기 위해 트래픽 도메인 토크나이저를 도입합니다. 교통 데이터. 자신의 데이터세트로 사용자 정의 토크나이저를 훈련하려면 코드에서 model_name
및 data_path
수정하세요.
model_name
: 네이티브 토크나이저가 포함된 기본 모델 경로입니다.
data_path
: 전처리 프로세스에서 추출된 학습 데이터 세트입니다.
코드를 사용하려면 명령을 따르십시오.
CD 토큰화 python Traffic_tokenizer.py
데이터 준비: 신경 언어 명령어 튜닝 데이터는 트래픽 분석 작업 이해를 위해 수집된 명령어 데이터 세트입니다.
튜닝 시작: 앞서 언급한 단계 후에 Trafficllm_stage1.sh를 사용하여 첫 번째 단계 튜닝을 시작할 수 있습니다. 아래와 같은 예가 있습니다:
PRE_SEQ_LEN=128 LR=2e-2 NUM_GPUS=1CUDA_VISIBLE_DEVICES=1 내보내기 torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py --do_train --train_file ../datasets/instructions/instructions.json --validation_file ../datasets/instructions/instructions.json --preprocessing_num_workers 10 --prompt_column 명령 --response_column 출력 --overwrite_cache --cache_dir ../cache --model_name_or_path ../models/chatglm2/chatglm2-6b --output_dir ../models/chatglm2/peft/instruction --overwrite_output_dir --max_source_length 1024 --max_target_length 32 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --gradient_accumulation_steps 16 --predict_with_generate --max_steps 20000 --logging_steps 10 --save_steps 4000 --learning_rate $LR --pre_seq_len $PRE_SEQ_LEN
데이터 준비: 작업별 트래픽 튜닝 데이터세트는 다양한 다운스트림 작업에 대한 전처리 단계에서 추출된 학습 데이터세트입니다.
튜닝 시작: 앞서 언급한 단계 후에 Trafficllm_stage2.sh를 사용하여 두 번째 단계 튜닝을 시작할 수 있습니다. 아래와 같은 예가 있습니다:
PRE_SEQ_LEN=128 LR=2e-2 NUM_GPUS=1CUDA_VISIBLE_DEVICES=1 내보내기 torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py --do_train --train_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_Detection_packet_train.json --validation_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_Detection_packet_train.json --preprocessing_num_workers 10 --prompt_column 명령 --response_column 출력 --overwrite_cache --cache_dir ../cache --model_name_or_path ../models/chatglm2/chatglm2-6b --output_dir ../models/chatglm2/peft/ustc-tfc-2016-Detection-packet --overwrite_output_dir --max_source_length 1024 --max_target_length 32 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --gradient_accumulation_steps 16 --predict_with_generate --max_steps 20000 --logging_steps 10 --save_steps 4000 --learning_rate $LR --pre_seq_len $PRE_SEQ_LEN
TrafficLLM은 EA-PEFT를 사용하여 확장 가능 적응을 통해 PEFT(매개변수 유효 미세 조정) 모델을 구성합니다. 이는 TrafficLLM이 쉽게 적응하는 데 도움이 됩니다. 새로운 환경. TrafficLLM 어댑터를 사용하면 이전 모델을 업데이트하거나 새 작업을 등록하는 유연한 작업이 가능합니다.
model_name
: 기본 모델의 경로입니다.
tuning_data
: 새 환경 데이터세트입니다.
adaptation_task
: 업데이트 또는 등록(기존 모델을 업데이트하거나 새 작업을 등록).
task_name
: 업데이트하거나 삽입할 다운스트림 작업 이름입니다.
MTD(맬웨어 트래픽 감지) 작업으로 TrafficLLM을 업데이트하는 예가 있습니다.
CD EA-PEFT python ea-peft.py --model_name /Your/Base/Model/Path --tuning_data /Your/New/Dataset/Path --adaptation_task update --task_name MTD
체크포인트: 자체 모델이나 출시된 체크포인트를 사용하여 TrafficLLM을 평가해 볼 수 있습니다.
데이터: 전처리 단계에서는 테스트 데이터 세트를 분할하고 평가를 위해 다양한 데이터 세트에 대한 레이블 파일을 구축합니다. 전처리 코드를 참고하세요.
다양한 다운스트림 작업에 대한 TrafficLLM의 효율성을 측정하려면 평가 코드를 실행하십시오.
model_name
: 기본 모델의 경로입니다.
traffic_task
: 탐지 작업 또는 생성 작업입니다.
test_file
: 전처리 단계에서 추출된 테스트 데이터세트입니다.
label_file
: 전처리 단계에서 추출된 라벨 파일입니다.
ptuning_path
: 작업별 평가를 위한 PEFT 모델 경로입니다.
MTD 작업에 대한 평가를 실행하는 예가 있습니다.
python 평가.py --model_name /Your/Base/Model/Path --traffic_task 감지 --test_file 데이터 세트/ustc-tfc-2016/ustc-tfc-2016_Detection_packet_test.json --label_file 데이터 세트/ustc-tfc-2016/ustc- tfc-2016_label.json --ptuning_path models/chatglm2/peft/ustc-tfc-2016-Detection-packet/checkpoints-20000/
로컬 장치에 TrafficLLM을 배포할 수 있습니다. 먼저, 훈련 단계에서 수집된 PEFT 모델을 등록하기 위해 config.json에서 모델 경로를 구성합니다. TrafficLLM에 6개의 작업을 등록하는 예가 있습니다.
{ "model_path": "models/chatglm2/chatglm2-6b/", "peft_path": "models/chatglm2/peft/", "peft_set": { "NLP": " 명령/체크포인트-8000/", "MTD": "ustc-tfc-2016-Detection-packet/checkpoint-10000/", "BND": "iscx-botnet-2014-Detection-packet/checkpoint-5000/", "WAD": "csic-2010-Detection-packet /checkpoint-6000/", "AAD": "dapt-2020-Detection-packet/checkpoint-20000/", "EVD": "iscx-vpn-2016-Detection-packet/checkpoint-4000/", "TBD": "iscx-tor-2016-Detection-packet /체크포인트-10000/" }, "tasks": { "맬웨어 트래픽 탐지": "MTD", "봇넷 탐지": "BND", "웹 공격 탐지": "WAD", "APT 공격 탐지": "AAD", "암호화된 VPN 탐지 ": "EVD", "Tor 동작 감지": "TBD" } }
그런 다음 inference.py 및 Trafficllm_server.py의 prepromt
함수에 preprompt를 추가해야 합니다. 사전 프롬프트는 작업별 트래픽 조정 중에 훈련 데이터에 사용되는 접두사 텍스트입니다.
터미널 모드에서 TrafficLLM과 채팅하려면 다음 명령을 실행할 수 있습니다:
python inference.py --config=config.json --prompt="Your Instruction Text +: + Traffic Data"
을 수행할 수 있습니다. 다음 명령을 사용하여 TrafficLLM의 웹 사이트 데모를 시작하십시오.
streamlit run Trafficllm_server.py
이 데모는 TrafficLLM의 웹 서버를 실행합니다. http://Your-Server-IP:Port
에 접속하여 채팅 상자에서 채팅하세요.
프레임워크와 코드의 기반이 되는 ChatGLM2 및 Llama2 관련 작업에 많은 감사를 드립니다. TrafficLLM 구축 디자인은 ET-BERT 및 GraphGPT에서 영감을 받았습니다. 그들의 훌륭한 작품에 감사드립니다.